Sort out some disagreements between our command line parsing and unit_test's defaults.

1 files changed, 66 insertions, 60 deletions

diff --git a/unit_tests.py b/unit_tests.py
index 9220609..b3c612e 100644
--- a/unit_tests.py
+++ b/unit_tests.py
@@ -10,65 +10,70 @@ from py11 import *
 from py11.mutex import MutexDB
 
 
-class Defaults:
-    so_pin    = "fnord"
-    user_pin  = "fnord"
-    only_slot = 0
-    verbose   = False
-    p11util   = "./libpkcs11.so"
-    dbname    = "unit_tests.db"
-    server    = "../libhal/tests/test-rpc_server"
-
-args = Defaults
+def main():
+    from sys import argv
+    global args
+    args = parse_arguments(argv[1:])
+    unittest.main(verbosity = 1 if args.quiet else 2, argv = argv[:1], catchbreak = True)
+
+def parse_arguments(argv = ()):
+    from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
+    parser = ArgumentParser(description = __doc__, formatter_class = ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--quiet",      action = "store_true",                          help = "suppress chatter")
+    parser.add_argument("--so-pin",     default = "fnord",                              help = "security officer PIN")
+    parser.add_argument("--user-pin",   default = "fnord",                              help = "user PIN")
+    parser.add_argument("--slot",       default = 0, type = int,                        help = "slot number")
+    parser.add_argument("--libpkcs11",  default = "./libpkcs11.so",                     help = "PKCS #11 library")
+    parser.add_argument("--p11util",    default = "./p11util",                          help = "p11util binary")
+    parser.add_argument("--dbname",     default = "unit_tests.db",                      help = "SQLite3 database")
+    parser.add_argument("--server",     default = "../libhal/tests/test-rpc_server",    help = "RPC server binary")
+    parser.add_argument("--all-tests",  action = "store_true",                          help = "enable tests usually skipped")
+    return parser.parse_args(argv)
+
+args = parse_arguments()
 p11  = None
 rpc  = None
 
-def main():
-    from argparse import ArgumentParser
-    parser = ArgumentParser(__doc__)
-    parser.add_argument("--verbose", action = "store_true")
-    parser.add_argument("--so-pin")
-    parser.add_argument("--user-pin")
-    parser.add_argument("--only-slot", type = int)
-    parser.add_argument("--p11util")
-    parser.add_argument("--dbname")
-    parser.add_argument("--server")
-    parser.set_defaults(**vars(Defaults))
-    global args
-    args = parser.parse_args()
-    unittest.main(verbosity = 2 if args.verbose else 1)
 
 def setUpModule():
-    global p11
-    p11 = PKCS11(args.p11util)
-
     from subprocess import Popen, PIPE
     from os import unlink, environ, geteuid
     from os.path import abspath, isfile
-
-    if args.verbose:
-        print "Initializing database"
-    db = abspath(args.dbname)
-    if isfile(db):
-        unlink(db)
-    environ["PKCS11_DATABASE"] = db
-    Popen((args.p11util, "-sup"), stdin = PIPE).communicate(
-        "{args.so_pin}\n{args.user_pin}\n".format(args = args))
+    global p11
+    global rpc
 
     if isfile(args.server):
-        if args.verbose:
-            print "Starting RPC server"
+        if not args.quiet:
+            print "Starting RPC server", args.server
         cmd = [args.server]
         if geteuid() != 0:
             cmd.insert(0, "sudo")
-        global rpc
+        if not args.quiet:
+            print "Running", " ".join(cmd)
         rpc = Popen(cmd)
 
-    if args.verbose:
+    if not args.quiet:
+        print "Setting PINs"
+    Popen((args.p11util, "-sup"), stdin = PIPE).communicate(
+        "{args.so_pin}\n{args.user_pin}\n".format(args = args))
+
+    if not args.quiet:
+        print "Configuring database", args.dbname
+    db = abspath(args.dbname)
+    if isfile(db):
+        unlink(db)
+    environ["PKCS11_DATABASE"] = db
+
+    if not args.quiet:
+        print "Loading PKCS #11 library", args.libpkcs11
+    p11 = PKCS11(args.libpkcs11)
+
+    if not args.quiet:
         print "Setup complete"
 
+
 def tearDownModule():
-    from os import unlink
+    from os import unlink, geteuid
     unlink(args.dbname)
     global rpc
     if rpc is not None:
@@ -116,19 +121,19 @@ class TestDevice(unittest.TestCase):
         p11.C_Finalize()
 
     def tearDown(self):
-        p11.C_CloseAllSessions(args.only_slot)
+        p11.C_CloseAllSessions(args.slot)
 
     def test_getSlots(self):
-        self.assertEqual(p11.C_GetSlotList(), (args.only_slot,))
+        self.assertEqual(p11.C_GetSlotList(), (args.slot,))
 
     def test_getTokenInfo(self):
-        token_info = p11.C_GetTokenInfo(args.only_slot)
+        token_info = p11.C_GetTokenInfo(args.slot)
         self.assertIsInstance(token_info, CK_TOKEN_INFO)
         self.assertEqual(token_info.label.rstrip(), "Cryptech Token")
 
     def test_sessions_serial(self):
-        rw_session = p11.C_OpenSession(args.only_slot, CKF_RW_SESSION | CKF_SERIAL_SESSION)
-        ro_session = p11.C_OpenSession(args.only_slot, CKF_SERIAL_SESSION)
+        rw_session = p11.C_OpenSession(args.slot, CKF_RW_SESSION | CKF_SERIAL_SESSION)
+        ro_session = p11.C_OpenSession(args.slot, CKF_SERIAL_SESSION)
 
     def test_sessions_parallel(self):
         # Cooked API doesn't allow this mistake, must use raw API to test
@@ -136,36 +141,36 @@ class TestDevice(unittest.TestCase):
         handle = CK_SESSION_HANDLE()
         notify = CK_NOTIFY()
         with self.assertRaises(CKR_SESSION_PARALLEL_NOT_SUPPORTED):
-            p11.so.C_OpenSession(args.only_slot, CKF_RW_SESSION, None, notify, byref(handle))
+            p11.so.C_OpenSession(args.slot, CKF_RW_SESSION, None, notify, byref(handle))
         with self.assertRaises(CKR_SESSION_PARALLEL_NOT_SUPPORTED):
-            p11.so.C_OpenSession(args.only_slot, 0,              None, notify, byref(handle))
+            p11.so.C_OpenSession(args.slot, 0,              None, notify, byref(handle))
 
     def test_login_user(self):
-        rw_session = p11.C_OpenSession(args.only_slot, CKF_RW_SESSION | CKF_SERIAL_SESSION)
-        ro_session = p11.C_OpenSession(args.only_slot, CKF_SERIAL_SESSION)
+        rw_session = p11.C_OpenSession(args.slot, CKF_RW_SESSION | CKF_SERIAL_SESSION)
+        ro_session = p11.C_OpenSession(args.slot, CKF_SERIAL_SESSION)
         p11.C_Login(ro_session, CKU_USER, args.user_pin)
         p11.C_Logout(ro_session)
 
     def test_login_so(self):
-        rw_session = p11.C_OpenSession(args.only_slot, CKF_RW_SESSION | CKF_SERIAL_SESSION)
-        ro_session = p11.C_OpenSession(args.only_slot, CKF_SERIAL_SESSION)
+        rw_session = p11.C_OpenSession(args.slot, CKF_RW_SESSION | CKF_SERIAL_SESSION)
+        ro_session = p11.C_OpenSession(args.slot, CKF_SERIAL_SESSION)
         self.assertRaises(CKR_SESSION_READ_ONLY_EXISTS, p11.C_Login, ro_session, CKU_SO, args.so_pin)
         p11.C_CloseSession(ro_session)
         p11.C_Login(rw_session, CKU_SO, args.so_pin)
-        self.assertRaises(CKR_SESSION_READ_WRITE_SO_EXISTS, p11.C_OpenSession, args.only_slot, CKF_SERIAL_SESSION)
+        self.assertRaises(CKR_SESSION_READ_WRITE_SO_EXISTS, p11.C_OpenSession, args.slot, CKF_SERIAL_SESSION)
         p11.C_Logout(rw_session)
 
     def test_random(self):
         # Testing that what this produces really is random seems beyond
         # the scope of a unit test.
-        session = p11.C_OpenSession(args.only_slot)
+        session = p11.C_OpenSession(args.slot)
         n = 17
         random = p11.C_GenerateRandom(session, n)
         self.assertIsInstance(random, str)
         self.assertEqual(len(random), n)
 
     def test_findObjects(self):
-        session = p11.C_OpenSession(args.only_slot)
+        session = p11.C_OpenSession(args.slot)
         p11.C_FindObjectsInit(session, CKA_CLASS = CKO_PUBLIC_KEY)
         with self.assertRaises(CKR_OPERATION_ACTIVE):
             p11.C_FindObjectsInit(session, CKA_CLASS = CKO_PRIVATE_KEY)
@@ -192,13 +197,13 @@ class TestKeys(unittest.TestCase):
         p11.C_Finalize()
 
     def setUp(self):
-        self.session = p11.C_OpenSession(args.only_slot)
+        self.session = p11.C_OpenSession(args.slot)
         p11.C_Login(self.session, CKU_USER, args.user_pin)
 
     def tearDown(self):
         for handle in p11.FindObjects(self.session):
             p11.C_DestroyObject(self.session, handle)
-        p11.C_CloseAllSessions(args.only_slot)
+        p11.C_CloseAllSessions(args.slot)
         del self.session
 
     def assertIsKeypair(self, public_handle, private_handle = None):
@@ -239,8 +244,8 @@ class TestKeys(unittest.TestCase):
         p11.C_VerifyInit(self.session, CKM_ECDSA_SHA256, public_key)
         p11.C_Verify(self.session, hamster, sig)
 
-    @unittest.skip("SHA-384 not available in current build")
     def test_gen_sign_verify_ecdsa_p384_sha384(self):
+        #if not args.all_tests: self.skipTest("SHA-384 not available in current build")
         public_key, private_key = p11.C_GenerateKeyPair(self.session, CKM_EC_KEY_PAIR_GEN,
                                                         CKA_ID = "EC-P384", CKA_EC_PARAMS = self.oid_p384,
                                                         CKA_SIGN = True, CKA_VERIFY = True)
@@ -252,8 +257,9 @@ class TestKeys(unittest.TestCase):
         p11.C_VerifyInit(self.session, CKM_ECDSA_SHA384, public_key)
         p11.C_Verify(self.session, hamster, sig)
 
-    @unittest.skip("SHA-512 not available in current build")
+
     def test_gen_sign_verify_ecdsa_p521_sha512(self):
+        #if not args.all_tests: self.skipTest("SHA-512 not available in current build")
         public_key, private_key = p11.C_GenerateKeyPair(self.session, CKM_EC_KEY_PAIR_GEN,
                                                         CKA_ID = "EC-P521", CKA_EC_PARAMS = self.oid_p521,
                                                         CKA_SIGN = True, CKA_VERIFY = True)
@@ -270,8 +276,8 @@ class TestKeys(unittest.TestCase):
           p11.C_GenerateKeyPair(self.session, CKM_RSA_PKCS_KEY_PAIR_GEN, CKA_MODULUS_BITS = 1024,
                                 CKA_ID = "RSA-1024", CKA_SIGN = True, CKA_VERIFY = True))
 
-    @unittest.skip("RSA key generation is still painfully slow")
     def test_gen_rsa_2048(self):
+        if not args.all_tests: self.skipTest("RSA key generation is still painfully slow")
         self.assertIsKeypair(
           p11.C_GenerateKeyPair(self.session, CKM_RSA_PKCS_KEY_PAIR_GEN, CKA_MODULUS_BITS = 2048,
                                 CKA_ID = "RSA-1024", CKA_SIGN = True, CKA_VERIFY = True))