Reworked AES implementation - use ctypes instead of m2crypto

Author: kamilbednarz

chef/aes.py

@@ -0,0 +1,134 @@
+import os
+
+from ctypes import *
+from rsa import load_crypto_lib, SSLError
+
+_eay = load_crypto_lib()
+
+c_int_p = POINTER(c_int)
+
+# void EVP_CIPHER_CTX_init(EVP_CIPHER_CTX *a);
+EVP_CIPHER_CTX_init = _eay.EVP_CIPHER_CTX_init
+EVP_CIPHER_CTX_init.argtypes = [c_void_p]
+EVP_CIPHER_CTX_init.restype = None
+
+#int EVP_CipherUpdate(EVP_CIPHER_CTX *ctx, unsigned char *out,
+#        int *outl, unsigned char *in, int inl);
+EVP_CipherUpdate = _eay.EVP_CipherUpdate
+EVP_CipherUpdate.argtypes = [c_void_p, c_char_p, c_int_p, c_char_p, c_int]
+EVP_CipherUpdate.restype = c_int
+
+#int EVP_CipherFinal(EVP_CIPHER_CTX *ctx, unsigned char *out,
+#         int *outl);
+EVP_CipherFinal = _eay.EVP_CipherFinal
+EVP_CipherFinal.argtypes = [c_void_p, c_char_p, c_int_p]
+EVP_CipherFinal.restype = c_int
+
+#EVP_CIPHER *EVP_aes_256_cbc(void);
+EVP_aes_256_cbc = _eay.EVP_aes_256_cbc
+EVP_aes_256_cbc.argtypes = []
+EVP_aes_256_cbc.restype = c_void_p
+
+#EVP_MD *EVP_sha1(void);
+EVP_sha1 = _eay.EVP_sha1
+EVP_sha1.argtypes = []
+EVP_sha1.restype = c_void_p
+
+#int EVP_CipherInit(EVP_CIPHER_CTX *ctx, const EVP_CIPHER *type,
+#        unsigned char *key, unsigned char *iv, int enc);
+EVP_CipherInit = _eay.EVP_CipherInit
+EVP_CipherInit.argtypes = [c_void_p, c_void_p, c_char_p, c_char_p, c_int]
+EVP_CipherInit.restype = c_int
+
+#int EVP_CIPHER_CTX_set_padding(EVP_CIPHER_CTX *x, int padding);
+EVP_CIPHER_CTX_set_padding = _eay.EVP_CIPHER_CTX_set_padding
+EVP_CIPHER_CTX_set_padding.argtypes = [c_void_p, c_int]
+EVP_CIPHER_CTX_set_padding.restype = c_int
+
+# Structures required for ctypes
+
+EVP_MAX_IV_LENGTH = 16
+EVP_MAX_BLOCK_LENGTH = 32
+AES_BLOCK_SIZE = 16
+
+class EVP_CIPHER(Structure):
+    _fields_ = [
+            ("nid", c_int),
+            ("block_size", c_int),
+            ("key_len", c_int),
+            ("iv_len", c_int),
+            ("flags", c_ulong),
+            ("init", c_voidp),
+            ("do_cipher", c_voidp),
+            ("cleanup", c_voidp),
+            ("set_asn1_parameters", c_voidp),
+            ("get_asn1_parameters", c_voidp),
+            ("ctrl", c_voidp),
+            ("app_data", c_voidp)
+    ]
+
+class EVP_CIPHER_CTX(Structure):
+    _fields_ = [
+            ("cipher", POINTER(EVP_CIPHER)),
+            ("engine", c_voidp),
+            ("encrypt", c_int),
+            ("buflen", c_int),
+            ("oiv", c_ubyte * EVP_MAX_IV_LENGTH),
+            ("iv", c_ubyte * EVP_MAX_IV_LENGTH),
+            ("buf", c_ubyte * EVP_MAX_BLOCK_LENGTH),
+            ("num", c_int),
+            ("app_data", c_voidp),
+            ("key_len", c_int),
+            ("flags", c_ulong),
+            ("cipher_data", c_voidp),
+            ("final_used", c_int),
+            ("block_mask", c_int),
+            ("final", c_ubyte * EVP_MAX_BLOCK_LENGTH) ]
+
+
+class AES256Cipher():
+    def __init__(self, key, iv, salt='12345678'):
+        self.key_data = create_string_buffer(key)
+        self.iv = create_string_buffer(iv)
+        self.encryptor = self.decryptor = None
+        self.salt = create_string_buffer(salt)
+
+        self.encryptor = EVP_CIPHER_CTX()
+        EVP_CIPHER_CTX_init(byref(self.encryptor))
+        EVP_CipherInit(byref(self.encryptor), EVP_aes_256_cbc(), self.key_data, self.iv, c_int(1))
+        EVP_CIPHER_CTX_set_padding(byref(self.encryptor), c_int(1))
+
+        self.decryptor = EVP_CIPHER_CTX()
+        EVP_CIPHER_CTX_init(byref(self.decryptor))
+        EVP_CipherInit(byref(self.decryptor), EVP_aes_256_cbc(), self.key_data, self.iv, c_int(0))
+        EVP_CIPHER_CTX_set_padding(byref(self.decryptor), c_int(1))
+
+    def encrypt(self, data):
+        length = c_int(len(data))
+        buf_length = c_int(length.value + AES_BLOCK_SIZE)
+        buf = create_string_buffer(buf_length.value)
+
+        final_buf = create_string_buffer(AES_BLOCK_SIZE)
+        final_length = c_int(0)
+
+        EVP_CipherUpdate(byref(self.encryptor), buf, byref(buf_length), create_string_buffer(data), length)
+        EVP_CipherFinal(byref(self.encryptor), final_buf, byref(final_length))
+
+        return string_at(buf, buf_length) + string_at(final_buf, final_length)
+
+    def decrypt(self, data):
+        length = c_int(len(data))
+        buf_length = c_int(length.value + AES_BLOCK_SIZE)
+        buf = create_string_buffer(buf_length.value)
+
+        final_buf = create_string_buffer(AES_BLOCK_SIZE)
+        final_length = c_int(0)
+
+        EVP_CipherUpdate(byref(self.decryptor), buf, byref(buf_length), create_string_buffer(data), length)
+        EVP_CipherFinal(byref(self.decryptor), final_buf, byref(final_length))
+
+        return string_at(buf, buf_length) + string_at(final_buf, final_length)
+
+
+
+

chef/encrypted_data_bag_item.py

@@ -58,7 +58,8 @@ def __init__(self, key, data):
             def encrypt(self):
                 if self.encrypted_data is None:
                     data = json.dumps({'json_wrapper': self.data})
-                    self.encrypted_data = self.encryptor.update(data) + self.encryptor.final()
+                    update_data = self.encryptor.update(data)
+                    self.encrypted_data = update_data + self.encryptor.final()
                     del self.encryptor
                 return self.encrypted_data
 

chef/tests/test_aes.py

@@ -0,0 +1,24 @@
+from chef.tests import ChefTestCase, TEST_ROOT
+from chef.aes import AES256Cipher
+from chef.rsa import SSLError
+
+import base64
+import os
+import hashlib
+import json
+
+class AES256CipherTestCase(ChefTestCase):
+    def setUp(self):
+        super(AES256CipherTestCase, self).setUp()
+        key = hashlib.sha256(open(os.path.join(TEST_ROOT, 'encryption_key')).read()).digest()
+        iv = base64.standard_b64decode('GLVikZLxG0SWYnb68Pr8Ag==\n')
+        self.cipher = AES256Cipher(key, iv)
+
+    def test_encrypt(self):
+        encrypted_value = self.cipher.encrypt('{"json_wrapper":"secr3t c0d3"}')
+        self.assertEquals(base64.standard_b64encode(encrypted_value).strip(), "Ym5T8umtSd0wgjDYq1ZDK5dAh6OjgrTxlloGNf2xYhg=")
+
+
+    def test_decrypt(self):
+        decrypted_value = self.cipher.decrypt(base64.standard_b64decode('Ym5T8umtSd0wgjDYq1ZDK5dAh6OjgrTxlloGNf2xYhg=\n'))
+        self.assertEquals(decrypted_value, '{"json_wrapper":"secr3t c0d3"}')