Use PyCrypto instead of Cryptlib to supply AES-ECB function.

Nothing wrong with Cryptlib, but we'll need an AES Keywrap
implementation to test key backup in our PyCrypto-based test harness,
so converting this implementation to PyCrypto is simplest overall.

1 files changed, 13 insertions, 19 deletions

diff --git a/aes_keywrap.py b/aes_keywrap.py
index 75aaa88..1d0be29 100644
--- a/aes_keywrap.py
+++ b/aes_keywrap.py
@@ -2,7 +2,7 @@
 
 """
 Python prototype of an AES Key Wrap implementation, RFC 5649 flavor
-per Russ, using Cryptlib to supply the AES code.
+per Russ, using PyCrypto to supply the AES code.
 """
 
 # Terminology mostly follows the RFC, including variable names.
@@ -15,9 +15,9 @@ per Russ, using Cryptlib to supply the AES code.
 # 64-bit blocks.
 
 
-from cryptlib_py import *
 from struct import pack, unpack
-import atexit
+from Crypto.Cipher import AES
+from array import array
 
 verbose = False
 
@@ -88,7 +88,7 @@ class Buffer(array):
 
 class KEK(object):
   """
-  Key encryption key, based on a Cryptlib encryption context.
+  Key encryption key, based on a PyCrypto encryption context.
 
   This can work with either Block objects or Python arrays.
 
@@ -98,10 +98,7 @@ class KEK(object):
   """
 
   def __init__(self, key):
-    self.ctx = cryptCreateContext(CRYPT_UNUSED, CRYPT_ALGO_AES)
-    atexit.register(cryptDestroyContext, self.ctx)
-    self.ctx.CTXINFO_MODE = CRYPT_MODE_ECB
-    self.ctx.CTXINFO_KEY  = key
+    self.ctx = AES.new(key, AES.MODE_ECB)
 
   def encrypt_block(self, i1, i2):
     """
@@ -109,9 +106,9 @@ class KEK(object):
     with AES-ECB, return the result split back into 64-bit blocks.
     """
 
-    aes_block = array("B", pack(">QQ", i1, i2))
-    cryptEncrypt(self.ctx, aes_block)
-    o1, o2 = tuple(Block(b) for b in unpack(">QQ", aes_block.tostring()))
+    aes_block = pack(">QQ", i1, i2)
+    aes_block = self.ctx.encrypt(aes_block)
+    o1, o2 = tuple(Block(b) for b in unpack(">QQ", aes_block))
     if verbose:
       print "  Encrypt: %s | %s  =>  %s | %s" % tuple(b.to_hex() for b in (i1, i2, o1, o2))
     return o1, o2
@@ -123,7 +120,7 @@ class KEK(object):
     """
 
     aes_block = b1 + b2
-    cryptEncrypt(self.ctx, aes_block)
+    aes_block = self.ctx.encrypt(aes_block.tostring())
     return Buffer(aes_block[:8]), Buffer(aes_block[8:])
 
   def decrypt_block(self, i1, i2):
@@ -132,9 +129,9 @@ class KEK(object):
     with AES-ECB, return the result split back into 64-bit blocks.
     """
 
-    aes_block = array("B", pack(">QQ", i1, i2))
-    cryptDecrypt(self.ctx, aes_block)
-    o1, o2 = tuple(Block(b) for b in unpack(">QQ", aes_block.tostring()))
+    aes_block = pack(">QQ", i1, i2)
+    aes_block = self.ctx.decrypt(aes_block)
+    o1, o2 = tuple(Block(b) for b in unpack(">QQ", aes_block))
     if verbose:
       print "  Decrypt: %s | %s  =>  %s | %s" % tuple(b.to_hex() for b in (i1, i2, o1, o2))
     return o1, o2
@@ -146,7 +143,7 @@ class KEK(object):
     """
 
     aes_block = b1 + b2
-    cryptDecrypt(self.ctx, aes_block)
+    aes_block = self.ctx.decrypt(aes_block.tostring())
     return Buffer(aes_block[:8]), Buffer(aes_block[8:])
 
 
@@ -506,9 +503,6 @@ if __name__ == "__main__":
   args = parser.parse_args()
   verbose = args.verbose
 
-  cryptInit()
-  atexit.register(cryptEnd)
-
   if args.under_test in ("long", "both"):
     print "Testing with Block (Python long) implementation"
     print