add --timeout

1 files changed, 19 insertions, 11 deletions

diff --git a/bin/cryptech_runcmd b/bin/cryptech_runcmd
index 6407043..61f277a 100755
--- a/bin/cryptech_runcmd
+++ b/bin/cryptech_runcmd
@@ -46,6 +46,7 @@ FPGA_CHUNK_SIZE = 4096
 default_pins = {'wheel': 'YouReallyNeedToChangeThisPINRightNowWeAreNotKidding',
                 'ct': 'ct',
                 }
+default_timeout = 2
 
 def parse_args():
     """
@@ -68,6 +69,13 @@ def parse_args():
                         help = "Username to use when logging into the HSM",
                         )
 
+    parser.add_argument("--timeout",
+                        metavar = 'SECONDS',
+                        type = int,
+                        default = default_timeout,
+                        help = "Timeout of commands (seconds)",
+                        )
+
     parser.add_argument('commands', metavar='CMD', type=str, nargs='+',
                         help='commands to execute')
     return parser.parse_args()
@@ -83,45 +91,45 @@ def _write(dst, data, debug=False):
         print("Wrote {!r}".format(data))
 
 
-def _read(dst, verbose=True):
+def _read(dst, args, verbose=True):
     res = ''
-    x = dst.read(1)
+    x = dst.read(args.timeout)
     while not x:
-        x = dst.read(1)
+        x = dst.read(args.timeout)
     while x:
         res += x
-        x = dst.read(1)
+        x = dst.read(args.timeout)
     #print ("Read {!r}".format(res))
     if verbose:
         sys.stdout.write(res)
     return res
 
-def _execute(dst, cmd):
+def _execute(dst, cmd, args):
     global default_pins
     _write(dst, '\r')
-    prompt = _read(dst)
+    prompt = _read(dst, args)
     if prompt.endswith('Username: '):
         _write(dst, args.username + "\r")
-        prompt = _read(dst)
+        prompt = _read(dst, args)
         if prompt.endswith('Password: '):
             pin = default_pins.get(args.username)
             if not pin:
                 pin = getpass.getpass("{} PIN: ".format(args.username))
             _write(dst, pin + '\r')
-            prompt = _read(dst)
+            prompt = _read(dst, args)
     if not prompt.endswith('> '):
         #sys.stderr.write('Device does not seem to be ready for a file transfer (got {!r})\n'.format(prompt))
         return prompt
     _write(dst, cmd + '\r')
-    response = _read(dst)
+    response = _read(dst, args)
     return response
 
 
 def main(args):
     global pin
-    dst = serial.Serial(args.device, 921600, timeout=2)
+    dst = serial.Serial(args.device, 921600, timeout=args.timeout)
     for this in args.commands:
-        _execute(dst, this)
+        _execute(dst, this, args)
     dst.close()
     return True