Clean up time-signature script.

1 files changed, 10 insertions, 3 deletions

diff --git a/scripts/time-signature.py b/scripts/time-signature.py
index 8b4cc5f..eef646d 100755
--- a/scripts/time-signature.py
+++ b/scripts/time-signature.py
@@ -14,7 +14,13 @@ from Crypto.Hash.SHA256 import SHA256Hash as SHA256
 from Crypto.Hash.SHA384 import SHA384Hash as SHA384
 from Crypto.Hash.SHA512 import SHA512Hash as SHA512
 
-from py11 import *
+try:
+    from py11 import *
+except ImportError:
+    # Kludge to let us run script from repo without installing py11
+    from os.path import dirname, abspath
+    sys.path.append(dirname(dirname(abspath(sys.argv[0]))))
+    from py11 import *
 
 
 if platform.system() == "Darwin":
@@ -90,7 +96,8 @@ def main():
     parser.add_argument("-s", "--slot",       default = 0, type = int,                  help = "slot number")
     parser.add_argument("-t", "--token",      action = "store_true",                    help = "store key on token")
     parser.add_argument("-l", "--libpkcs11",  default = libpkcs11_default,              help = "PKCS #11 library")
-    parser.add_argument("-k", "--key",        choices = tuple(key_table), nargs = 1,    help = "test just one key")
+    parser.add_argument("-k", "--keys",       choices = tuple(key_table), nargs = "+",
+                                              default = list(key_table),                help = "keys to test")
     parser.add_argument("-q", "--quiet",      action = "store_true",                    help = "be less chatty")
     args = parser.parse_args()
 
@@ -103,7 +110,7 @@ def main():
     session = p11.C_OpenSession(args.slot)
     p11.C_Login(session, CKU_USER, args.pin)
 
-    for name in args.key or key_table:
+    for name in args.keys:
 
         print "Starting test with key {}, {} iterations".format(name, args.iterations)