Clean up test_hashes script.

1 files changed, 25 insertions, 35 deletions

diff --git a/tests/test_hashes.py b/tests/test_hashes.py
index 2e38ca3..d84fe6b 100644
--- a/tests/test_hashes.py
+++ b/tests/test_hashes.py
@@ -12,8 +12,8 @@ atexit.register(cryptDeviceClose, hwdev)
 
 # Usual NIST sample messages.
 
-def hextext(s):
-    return "".join(s.split()).lower()
+def hextext(text):
+    return "".join(text.split()).lower()
 
 NIST_512_SINGLE      = "abc"
 SHA1_SINGLE_DIGEST   = hextext("A9993E36 4706816A BA3E2571 7850C26C 9CD0D89D")
@@ -36,44 +36,34 @@ SHA384_DOUBLE_DIGEST = hextext("09330C33 F71147E8 3D192FC7 82CD1B47 53111B17 3B3
 SHA512_DOUBLE_DIGEST = hextext("8E959B75 DAE313DA 8CF4F728 14FC143F 8F7779C6 EB9F7FA1 7299AEAD B6889018"
                                "501D289E 4900F7E4 331B99DE C4B5433A C7D329EE B6DD2654 5E96E55B 874BE909")
 
-def do_hash(ctx, s):
+def do_hash(dev, algorithm, text, blocksize = None):
+    ctx = None
     try:
-        cryptEncrypt(ctx, array("c", s))
+        if dev is None:
+            ctx = cryptCreateContext(CRYPT_UNUSED, algorithm)
+        else:
+            ctx = cryptDeviceCreateContext(dev, algorithm)
+        if blocksize is not None:
+            ctx.CTXINFO_BLOCKSIZE = blocksize
+        cryptEncrypt(ctx, array("c", text))
         cryptEncrypt(ctx, array("c", ""))
         result = ctx.CRYPT_CTXINFO_HASHVALUE
         return result.encode("hex")
     finally:
-        cryptDestroyContext(ctx)
-
-def sha1(d, s):
-    if d is None:
-        ctx = cryptCreateContext(CRYPT_UNUSED, CRYPT_ALGO_SHA1)
-    else:
-        ctx = cryptDeviceCreateContext(d, CRYPT_ALGO_SHA1)
-    return do_hash(ctx, s)
-
-def sha256(d, s):
-    if d is None:
-        ctx = cryptCreateContext(CRYPT_UNUSED, CRYPT_ALGO_SHA2)
-    else:
-        ctx = cryptDeviceCreateContext(d, CRYPT_ALGO_SHA2)
-    return do_hash(ctx, s)
-
-def sha384(d, s):
-    if d is None:
-        ctx = cryptCreateContext(CRYPT_UNUSED, CRYPT_ALGO_SHA2)
-    else:
-        ctx = cryptDeviceCreateContext(d, CRYPT_ALGO_SHA2)
-    ctx.CTXINFO_BLOCKSIZE = 48
-    return do_hash(ctx, s)
-
-def sha512(d, s):
-    if d is None:
-        ctx = cryptCreateContext(CRYPT_UNUSED, CRYPT_ALGO_SHA2)
-    else:
-        ctx = cryptDeviceCreateContext(d, CRYPT_ALGO_SHA2)
-    ctx.CTXINFO_BLOCKSIZE = 64
-    return do_hash(ctx, s)
+        if ctx is not None:
+            cryptDestroyContext(ctx)
+
+def sha1(dev, text):
+    return do_hash(dev, CRYPT_ALGO_SHA1, text)
+
+def sha256(dev, text):
+    return do_hash(dev, CRYPT_ALGO_SHA2, text)
+
+def sha384(dev, text):
+    return do_hash(dev, CRYPT_ALGO_SHA2, text, 48)
+
+def sha512(dev, text):
+    return do_hash(dev, CRYPT_ALGO_SHA2, text, 64)
 
 have_i2c = os.path.exists("/dev/i2c-2")