Merge pull request #386 from t184256/aes-fixes

AES fixes

Author: tomato42

tlslite/utils/openssl_aes.py

@@ -9,44 +9,52 @@
 if m2cryptoLoaded:
 
     def new(key, mode, IV):
+        # IV argument name is a part of the interface
+        # pylint: disable=invalid-name
         return OpenSSL_AES(key, mode, IV)
 
     class OpenSSL_AES(AES):
 
         def __init__(self, key, mode, IV):
+            # IV argument/field names are a part of the interface
+            # pylint: disable=invalid-name
             AES.__init__(self, key, mode, IV, "openssl")
-            self.key = key
-            self.IV = IV
+            self.IV, self._key = IV, key
+            self._context = None
+            self._encrypt = None
 
-        def _createContext(self, encrypt):
-            context = m2.cipher_ctx_new()
-            if len(self.key)==16:
+        def _init_context(self, encrypt=True):
+            if len(self._key) == 16:
                 cipherType = m2.aes_128_cbc()
-            if len(self.key)==24:
+            if len(self._key) == 24:
                 cipherType = m2.aes_192_cbc()
-            if len(self.key)==32:
+            if len(self._key) == 32:
                 cipherType = m2.aes_256_cbc()
-            m2.cipher_init(context, cipherType, self.key, self.IV, encrypt)
-            return context
+            self._context = m2.cipher_ctx_new()
+            m2.cipher_init(self._context, cipherType, self._key, self.IV,
+                           int(encrypt))
+            m2.cipher_set_padding(self._context, 0)
+            self._encrypt = encrypt
 
         def encrypt(self, plaintext):
+            if self._context is None:
+                self._init_context(encrypt=True)
+            else:
+                assert self._encrypt, '.encrypt() not allowed after .decrypt()'
             AES.encrypt(self, plaintext)
-            context = self._createContext(1)
-            ciphertext = m2.cipher_update(context, plaintext)
-            m2.cipher_ctx_free(context)
-            self.IV = ciphertext[-self.block_size:]
+            ciphertext = m2.cipher_update(self._context, plaintext)
             return bytearray(ciphertext)
 
         def decrypt(self, ciphertext):
+            if self._context is None:
+                self._init_context(encrypt=False)
+            else:
+                assert not self._encrypt, \
+                       '.decrypt() not allowed after .encrypt()'
             AES.decrypt(self, ciphertext)
-            context = self._createContext(0)
-            #I think M2Crypto has a bug - it fails to decrypt and return the last block passed in.
-            #To work around this, we append sixteen zeros to the string, below:
-            plaintext = m2.cipher_update(context, ciphertext+(b'\0'*16))
-
-            #If this bug is ever fixed, then plaintext will end up having a garbage
-            #plaintext block on the end.  That's okay - the below code will discard it.
-            plaintext = plaintext[:len(ciphertext)]
-            m2.cipher_ctx_free(context)
-            self.IV = ciphertext[-self.block_size:]
+            plaintext = m2.cipher_update(self._context, ciphertext)
             return bytearray(plaintext)
+
+        def __del__(self):
+            if self._context is not None:
+                m2.cipher_ctx_free(self._context)

tlslite/utils/python_aes.py

@@ -11,19 +11,24 @@
 
 
 def new(key, mode, IV):
+    # IV argument name is a part of the interface
+    # pylint: disable=invalid-name
     return Python_AES(key, mode, IV)
 
 
 class Python_AES(AES):
     def __init__(self, key, mode, IV):
+        # IV argument/field names are a part of the interface
+        # pylint: disable=invalid-name
+        key, IV = bytearray(key), bytearray(IV)
         super(Python_AES, self).__init__(key, mode, IV, "python")
         self.rijndael = Rijndael(key, 16)
         self.IV = IV
 
     def encrypt(self, plaintext):
         super(Python_AES, self).encrypt(plaintext)
 
-        plaintextBytes = plaintext[:]
+        plaintextBytes = bytearray(plaintext)
         chainBytes = self.IV[:]
 
         #CBC Mode: For each block...

unit_tests/test_tlslite_utils_aes_split.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2019, Alexander Sosedkin
+#
+# See the LICENSE file for legal information regarding use of this file.
+
+# compatibility with Python 2.6, for that we need unittest2 package,
+# which is not available on 3.3 or 3.4
+try:
+    import unittest2 as unittest
+except ImportError:
+    import unittest
+
+import sys
+
+from hypothesis import given, assume, settings
+from hypothesis.strategies import binary, integers, tuples
+
+from tlslite.utils import cryptomath
+
+import tlslite.utils.python_aes
+py_aes = lambda key, iv: tlslite.utils.python_aes.new(key, 2, iv)
+
+
+HYP_SETTINGS = {'deadline': None} if sys.version_info > (2, 7) else {}
+
+
+class TestAES(unittest.TestCase):
+    _given = given(binary(min_size=24, max_size=24),          # key
+                   binary(min_size=16, max_size=16),          # iv
+                   binary(min_size=13*16, max_size=13*16),    # plaintext
+                   (tuples(integers(0, 13), integers(0, 13))  # split points
+                       .filter(lambda split_pts: split_pts[0] <= split_pts[1])
+                       .map(lambda lengths: [i * 16 for i in lengths])))
+
+    def split_test(self, key, iv, plaintext, split_points, make_impl=py_aes):
+        i, j = split_points
+
+        ciphertext = make_impl(key, iv).encrypt(plaintext)
+        self.assertEqual(make_impl(key, iv).decrypt(ciphertext), plaintext)
+
+        impl = make_impl(key, iv)
+        pl1, pl2, pl3 = plaintext[:i], plaintext[i:j], plaintext[j:]
+        ci1, ci2, ci3 = impl.encrypt(pl1), impl.encrypt(pl2), impl.encrypt(pl3)
+        self.assertEqual(ci1 + ci2 + ci3, ciphertext)
+
+        impl = make_impl(key, iv)
+        pl1, pl2, pl3 = impl.decrypt(ci1), impl.decrypt(ci2), impl.decrypt(ci3)
+        self.assertEqual(pl1 + pl2 + pl3, plaintext)
+
+        return ciphertext
+
+    @_given
+    @settings(**HYP_SETTINGS)
+    def test_python(self, key, iv, plaintext, split_points):
+        self.split_test(key, iv, plaintext, split_points)
+
+    @unittest.skipIf(not cryptomath.m2cryptoLoaded, "requires M2Crypto")
+    @_given
+    @settings(**HYP_SETTINGS)
+    def test_python_vs_mcrypto(self, key, iv, plaintext, split_points):
+        import tlslite.utils.openssl_aes
+        m2_aes = lambda k, iv: tlslite.utils.openssl_aes.new(k, 2, iv)
+
+        py_res = self.split_test(key, iv, plaintext, split_points, py_aes)
+        m2_res = self.split_test(key, iv, plaintext, split_points, m2_aes)
+        self.assertEqual(py_res, m2_res)
+
+    @unittest.skipIf(not cryptomath.pycryptoLoaded, "requires pycrypto")
+    @_given
+    @settings(**HYP_SETTINGS)
+    def test_python_vs_pycrypto(self, key, iv, plaintext, split_points):
+        import tlslite.utils.pycrypto_aes
+        pc_aes = lambda k, iv: tlslite.utils.pycrypto_aes.new(k, 2, iv)
+
+        py_res = self.split_test(key, iv, plaintext, split_points, py_aes)
+        pc_res = self.split_test(key, iv, plaintext, split_points, pc_aes)
+        self.assertEqual(py_res, pc_res)