Add scaling factor to trng_extractor for easier bulk data extraction.

1 files changed, 26 insertions, 14 deletions

diff --git a/i2c/sw/trng_extractor_i2c.c b/i2c/sw/trng_extractor_i2c.c
index 29f054c..7b5b3b2 100644
--- a/i2c/sw/trng_extractor_i2c.c
+++ b/i2c/sw/trng_extractor_i2c.c
@@ -38,6 +38,7 @@
 #include <stdlib.h>
 #include <unistd.h>
 #include <stdint.h>
+#include <ctype.h>
 
 #include "tc_i2c.h"
 #include "cryptech_memory_map.h"
@@ -47,8 +48,8 @@ char *usage =
 \n\
 -a      avalanche entropy\n\
 -r      rosc entropy\n\
--c      csprng\n\
--n      number of 4-byte samples (defaults to 10)\n\
+-c      csprng (default data source)\n\
+-n      number of 4-byte samples (scale with K, M, or G suffix)\n\
 -o      output file (defaults to stdout)\n\
 ";
 
@@ -73,10 +74,13 @@ static int extract(off_t status_addr, off_t data_addr, uint32_t *data)
 /* main */
 int main(int argc, char *argv[])
 {
-    int i, opt, num_words = 10;
-    uint32_t data;
-    off_t status_addr = 0, data_addr = 0;
+    int i, opt;
+    unsigned long num_words = 1;
+    char *endptr;
+    off_t status_addr = CSPRNG_ADDR_STATUS;
+    off_t data_addr = CSPRNG_ADDR_RANDOM;
     FILE *output = stdout;
+    uint32_t data;
     char *dev = I2C_dev;
     int addr = I2C_addr;
 
@@ -100,11 +104,23 @@ int main(int argc, char *argv[])
             data_addr = CSPRNG_ADDR_RANDOM;
             break;
         case 'n':
-            num_words = atoi(optarg);
-            if (num_words <= 0) {
-                fprintf(stderr, "-n requires a positive integer argument\n");
-                return EXIT_FAILURE;
-            }
+	    num_words = strtoul(optarg, &endptr, 10);
+	    switch (toupper(*endptr)) {
+	    case '\0':
+		break;
+	    case 'K':
+		num_words *= 1000;
+		break;
+	    case 'M':
+		num_words *= 1000000;
+		break;
+	    case 'G':
+		num_words *= 1000000000;
+		break;
+	    default:
+		fprintf(stderr, "unsupported -n suffix %s\n", endptr);
+		return EXIT_FAILURE;
+	    }
             break;
         case 'o':
             output = fopen(optarg, "wb+");
@@ -140,10 +156,6 @@ int main(int argc, char *argv[])
         fprintf(stderr, "\n");
         goto errout;
     }
-    if (status_addr == 0) {
-        fprintf(stderr, "%s: a data source must be specified\n", argv[0]);
-        goto errout;
-    }
 
     /* set up I2C */
     if (i2c_open(dev, addr) != 0) {