add test scriptHEADmaster

1 files changed, 30 insertions, 14 deletions

diff --git a/bin/cryptech_runcmd b/bin/cryptech_runcmd
index 61f277a..b586c16 100755
--- a/bin/cryptech_runcmd
+++ b/bin/cryptech_runcmd
@@ -91,43 +91,59 @@ def _write(dst, data, debug=False):
         print("Wrote {!r}".format(data))
 
 
-def _read(dst, args, verbose=True):
+def _read(dst, retry=10, verbose=True, debug=False):
     res = ''
-    x = dst.read(args.timeout)
-    while not x:
-        x = dst.read(args.timeout)
+    x = dst.read(1)
+    while not x and retry:
+        x = dst.read(1)
+        retry -= 1
     while x:
         res += x
-        x = dst.read(args.timeout)
-    #print ("Read {!r}".format(res))
+        x = dst.read(1)
+    if debug:
+        print ("Read {!r}".format(res))
     if verbose:
         sys.stdout.write(res)
     return res
 
+def _read_until_prompt(dst, retry):
+    _write(dst, '\r')
+    while retry:
+        prompt = _read(dst, retry = 1)
+        retry -= 1
+        if prompt.endswith('Username: ') or \
+                prompt.endswith('Password: ') or \
+                prompt.endswith('> '):
+            return prompt
+    return ''
+
 def _execute(dst, cmd, args):
     global default_pins
-    _write(dst, '\r')
-    prompt = _read(dst, args)
+
+    prompt = _read_until_prompt(dst, retry = args.timeout)
+
     if prompt.endswith('Username: '):
-        _write(dst, args.username + "\r")
-        prompt = _read(dst, args)
+        _write(dst, args.username + "\r", debug = True)
+        prompt = _read(dst, retry = args.timeout)
         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, args)
+            prompt = _read(dst, retry = args.timeout)
     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, args)
+
+    _write(dst, cmd)
+
+    response = _read_until_prompt(dst, retry = args.timeout)
     return response
 
 
 def main(args):
     global pin
-    dst = serial.Serial(args.device, 921600, timeout=args.timeout)
+    dst = serial.Serial(args.device, 921600, timeout=1)
     for this in args.commands:
         _execute(dst, this, args)
     dst.close()