Merge pull request #33 from danielinux/aes-ctr

Author: haydenroche5

tests/test_ciphers.py

@@ -51,14 +51,15 @@
     from wolfcrypt.ciphers import (Ed448Private, Ed448Public)
 
 from wolfcrypt.ciphers import (
-    MODE_ECB, MODE_CBC, WolfCryptError
+    MODE_CTR, MODE_ECB, MODE_CBC, WolfCryptError
 )
 
 
 @pytest.fixture
 def vectors():
-    TestVector = namedtuple("TestVector", """key iv plaintext ciphertext raw_key
-                                             pkcs8_key pem""")
+    TestVector = namedtuple("TestVector", """key iv plaintext ciphertext 
+                ciphertext_ctr raw_key
+                pkcs8_key pem""")
     TestVector.__new__.__defaults__ = (None,) * len(TestVector._fields)
 
     # test vector dictionary
@@ -69,7 +70,8 @@ def vectors():
             key="0123456789abcdef",
             iv="1234567890abcdef",
             plaintext=t2b("now is the time "),
-            ciphertext=h2b("959492575f4281532ccc9d4677a233cb")
+            ciphertext=h2b("959492575f4281532ccc9d4677a233cb"),
+            ciphertext_ctr = h2b('287528ddf484b1055debbe751eb52b8a')
         )
     if _lib.CHACHA_ENABLED:
         vectorArray[ChaCha]=TestVector(
@@ -235,12 +237,12 @@ def cipher_new(cipher_cls, vectors):
         MODE_CBC,
         vectors[cipher_cls].iv)
 
-
 def test_block_cipher(cipher_cls, vectors):
     key = vectors[cipher_cls].key
     iv = vectors[cipher_cls].iv
     plaintext = vectors[cipher_cls].plaintext
     ciphertext = vectors[cipher_cls].ciphertext
+    ciphertext_ctr = vectors[cipher_cls].ciphertext_ctr
 
     with pytest.raises(ValueError):
         cipher_cls.new(key[:-1], MODE_CBC, iv)  # invalid key length
@@ -257,11 +259,21 @@ def test_block_cipher(cipher_cls, vectors):
     with pytest.raises(ValueError):
         cipher_cls.new(key, MODE_CBC, iv[:-1])  # invalid iv length
 
+
+    # Test AES in counter mode
+    if ciphertext_ctr is not None:
+        cipher_obj = cipher_cls.new(key, MODE_CTR, iv)
+        res = cipher_obj.encrypt(plaintext)
+        assert res == ciphertext_ctr
+        cipher_obj = cipher_cls.new(key, MODE_CTR, iv)
+        assert plaintext == cipher_obj.decrypt(res)
+
     # single encryption
     cipher_obj = cipher_new(cipher_cls, vectors)
 
     assert cipher_obj.encrypt(plaintext) == ciphertext
 
+
     # many encryptions
     cipher_obj = cipher_new(cipher_cls, vectors)
     result = t2b("")

wolfcrypt/_build_ffi.py

@@ -324,6 +324,7 @@ def generate_libwolfssl():
     int wc_AesSetKey(Aes*, const byte*, word32, const byte*, int);
     int wc_AesCbcEncrypt(Aes*, byte*, const byte*, word32);
     int wc_AesCbcDecrypt(Aes*, byte*, const byte*, word32);
+    int wc_AesCtrEncrypt(Aes*, byte*, const byte*, word32);
     """
 
 if AES_ENABLED and AESGCM_STREAM:

wolfcrypt/_build_wolfssl.py

@@ -184,6 +184,7 @@ def make_flags(prefix):
 
         # symmetric ciphers
         flags.append("--enable-aes")
+        flags.append("--enable-aesctr")
         flags.append("--enable-des3")
         flags.append("--enable-chacha")
 

wolfcrypt/ciphers.py

@@ -92,12 +92,14 @@ def __init__(self, key, mode, IV=None):
         if mode not in _FEEDBACK_MODES:
             raise ValueError("this mode is not supported")
 
-        if mode == MODE_CBC:
+        if mode == MODE_CBC or mode == MODE_CTR:
             if IV is None:
                 raise ValueError("this mode requires an 'IV' string")
         else:
             raise ValueError("this mode is not supported by this cipher")
 
+        self.mode = mode
+
         if self.key_size:
             if self.key_size != len(key):
                 raise ValueError("key must be %d in length, not %d" %
@@ -218,17 +220,31 @@ def _set_key(self, direction):
             if direction == _ENCRYPTION:
                 return _lib.wc_AesSetKey(
                     self._enc, self._key, len(self._key), self._IV, _ENCRYPTION)
-
+            if self.mode == MODE_CTR:
+                return _lib.wc_AesSetKey(
+                    self._dec, self._key, len(self._key), self._IV, _ENCRYPTION)
             return _lib.wc_AesSetKey(
                 self._dec, self._key, len(self._key), self._IV, _DECRYPTION)
 
         def _encrypt(self, destination, source):
-            return _lib.wc_AesCbcEncrypt(self._enc, destination,
-                                         source, len(source))
+            if self.mode == MODE_CBC:
+                return _lib.wc_AesCbcEncrypt(self._enc, destination,
+                        source, len(source))
+            elif self.mode == MODE_CTR:
+                return _lib.wc_AesCtrEncrypt(self._enc, destination,
+                        source, len(source))
+            else:
+                raise ValueError("Invalid mode associated to cipher")
 
         def _decrypt(self, destination, source):
-            return _lib.wc_AesCbcDecrypt(self._dec, destination,
-                                         source, len(source))
+            if self.mode == MODE_CBC:
+                return _lib.wc_AesCbcDecrypt(self._dec, destination,
+                        source, len(source))
+            elif self.mode == MODE_CTR:
+                return _lib.wc_AesCtrEncrypt(self._dec, destination,
+                        source, len(source))
+            else:
+                raise ValueError("Invalid mode associated to cipher")
 
 if _lib.AESGCM_STREAM:
     class _AesGcmStream(object):