Add aes 192 256 support (#16)

* added support for 192 and 256 bit keys. added nist tests. did not add aes core tests

* adeed core specic tests

Author: bjohnsonfl

src/python/aes_decrypt.py

@@ -2,10 +2,8 @@
 
 logger = logging.getLogger(__name__)
 
-# AES128 Constants
-Nr = 10
+# AES Constants
 Nb = 4
-Nk = 4
 
 s_box = [[0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76], 
          [0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0], 
@@ -149,14 +147,28 @@ def InvCipher(text_in, w):
     if isinstance(text_in_hex, str):
         text_in_hex = bytes.fromhex(text_in)
     state = []
+
+    key_len = len(w) * 4
+    if key_len == 128:
+        Nk = 4
+        Nr = 10
+    elif key_len == 192:
+        Nk = 6
+        Nr = 12
+    elif key_len == 256:
+        Nk = 8
+        Nr = 14
+    else:
+        assert(0)
+
     # s[r,c]=in[r+4c] for0≤r<4 and 0≤c<Nb, (3.3)
     for r in range(0,4):
         row = []
         for c in range(0,Nb):
             row.append(text_in_hex[r+4*c])
         state.append(row)
 
-    w = KeyExpansion(w, Nk)
+    w = KeyExpansion(w, Nk, Nr)
 
     logger.debug("\nRound: 0")
     logger.debug("iinput")
@@ -181,7 +193,7 @@ def InvCipher(text_in, w):
     logger.debug(f'out: {out}')
     return out
 
-def KeyExpansion(key, Nk):
+def KeyExpansion(key, Nk, Nr):
     temp = None
     w = []
     key_hex = bytes.fromhex(key)
@@ -201,8 +213,10 @@ def KeyExpansion(key, Nk):
 
     for i in range(Nk, Nb * (Nr+1)):
         temp = w[i-1]
-        if (i%Nk ==0):
+        if (i%Nk==0):
             temp = SubWord(RotWord(temp)) ^ rcon_arr[i//Nk]
+        elif (Nk > 6 and (i % Nk) == 4):
+            temp = SubWord(temp)
         w.append( w[i-Nk] ^ temp )
         logger.debug(f'{i//4}: {hex(w[i])}')
 

src/python/aes_encrypt.py

@@ -2,10 +2,8 @@
 
 logger = logging.getLogger(__name__)
 
-# AES128 Constants
-Nr = 10
+# AES Constants
 Nb = 4
-Nk = 4
 
 s_box = [[0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76], 
          [0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0], 
@@ -134,6 +132,20 @@ def Cipher(text_in, w):
         logger.info("input was string not binary")
         text_in_hex = bytes.fromhex(text_in)
     state = []
+
+    key_len = len(w) * 4
+    if key_len == 128:
+        Nk = 4
+        Nr = 10
+    elif key_len == 192:
+        Nk = 6
+        Nr = 12
+    elif key_len == 256:
+        Nk = 8
+        Nr = 14
+    else:
+        assert(0)
+
     # s[r,c]=in[r+4c] for0≤r<4 and 0≤c<Nb, (3.3)
     for r in range(0,4):
         row = []
@@ -142,7 +154,7 @@ def Cipher(text_in, w):
         state.append(row)
     print_state(state)
 
-    w = KeyExpansion(w, Nk)
+    w = KeyExpansion(w, Nk, Nr)
 
     logger.debug("\nRound: 0")
     logger.debug("input")
@@ -168,7 +180,7 @@ def Cipher(text_in, w):
     logger.debug(f'out: {out}')
     return out
 
-def KeyExpansion(key, Nk):
+def KeyExpansion(key, Nk, Nr):
     temp = None
     w = []
     key_hex = bytes.fromhex(key)
@@ -187,8 +199,10 @@ def KeyExpansion(key, Nk):
 
     for i in range(Nk, Nb * (Nr+1)):
         temp = w[i-1]
-        if (i%Nk ==0):
+        if (i%Nk == 0):
             temp = SubWord(RotWord(temp)) ^ rcon_arr[i//Nk]
+        elif (Nk > 6 and (i % Nk) == 4):
+            temp = SubWord(temp)
         w.append( w[i-Nk] ^ temp )
         logger.debug(f'{i//4}: {hex(w[i])}')
 

test/nist/aesmmt/test_aes_cbc.py

@@ -4,9 +4,8 @@
 from src.python.aes_modes import aes_cbc
 from scripts.rsp_parser import parse_kat_file
 
-# TODO: Support 192 and 256
 def get_mmt_cbc_files():
-    return list(Path("nist/aesmmt").glob("CBC*128.rsp"))
+    return list(Path("nist/aesmmt").glob("CBC*.rsp"))
 
 @pytest.mark.parametrize("mmt_files", get_mmt_cbc_files(), ids=lambda p: p.name)
 def test_mmt_aes_cbc_encrypt(mmt_files):

test/nist/aesmmt/test_aes_ecb.py

@@ -4,15 +4,15 @@
 from src.python.aes_modes import aes_ecb
 from scripts.rsp_parser import parse_kat_file
 
-# TODO: Support 192 and 256
 def get_mmt_ecb_files():
-    return list(Path("nist/aesmmt").glob("ECB*128.rsp"))
+    return list(Path("nist/aesmmt").glob("ECB*.rsp"))
 
 @pytest.mark.parametrize("mmt_files", get_mmt_ecb_files(), ids=lambda p: p.name)
 def test_mmt_aes_ecb_encrypt(mmt_files):
     mmt_list = parse_kat_file(mmt_files)
     for (op, count), value in mmt_list.items():
         if op == "ENCRYPT":
+            print(f"ENCRYPT: {count}")
             ct = aes_ecb(value.plaintext, value.key, op)
             assert(ct == value.ciphertext)
 
@@ -21,5 +21,6 @@ def test_mmt_aes_ecb_decrypt(mmt_files):
     mmt_list = parse_kat_file(mmt_files)
     for (op, count), value in mmt_list.items():
         if op == "DECRYPT":
+            print(f"DECRYPT: {count}")
             pt = aes_ecb(value.ciphertext, value.key, op)
             assert(pt == value.plaintext)

test/python/test_aes_decrypt.py

@@ -7,6 +7,16 @@ def test_decrypt_cipher():
     out1 = "00112233445566778899aabbccddeeff"
     assert(InvCipher(input1, key1) == out1)
 
+    input_192 = "dda97ca4864cdfe06eaf70a0ec0d7191"
+    key_192 = "000102030405060708090a0b0c0d0e0f1011121314151617"
+    output_192 = "00112233445566778899aabbccddeeff"
+    assert(InvCipher(input_192, key_192) == output_192)
+
+    input_256 = "8ea2b7ca516745bfeafc49904b496089"
+    key_256 = "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f"
+    output_256 = "00112233445566778899aabbccddeeff"
+    assert(InvCipher(input_256, key_256) == output_256)
+
 if __name__ == "__main__":
     logging.basicConfig(
         level=logging.DEBUG,

test/python/test_aes_encrypt.py

@@ -1,9 +1,24 @@
 import logging
-from src.python.aes_encrypt import Cipher, KeyExpansion, Nk
+from src.python.aes_encrypt import Cipher, KeyExpansion
 
 def test_cipher():
-    key = "2b7e151628aed2a6abf7158809cf4f3c"
-    KeyExpansion(key, Nk)
+    keys = ["2b7e151628aed2a6abf7158809cf4f3c",
+           "8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b",
+           "603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4"]
+    for key in keys:
+        key_len = len(key) * 4
+        if key_len == 128:
+            Nk = 4
+            Nr = 10
+        elif key_len == 192:
+            Nk = 6
+            Nr = 12
+        elif key_len == 256:
+            Nk = 8
+            Nr = 14
+        else:
+            assert(0)
+        KeyExpansion(key, Nk, Nr)
 
     # Input      = 32 43 f6 a8 88 5a 30 8d 31 31 98 a2 e0 37 07 34 
     # Cipher Key = 2b 7e 15 16 28 ae d2 a6 ab f7 15 88 09 cf 4f 3c
@@ -17,6 +32,16 @@ def test_cipher():
     out2 = "69c4e0d86a7b0430d8cdb78070b4c55a" # Appendix C
     assert(Cipher(input2, key2) == out2)
 
+    input_192 = "00112233445566778899aabbccddeeff"
+    key_192 = "000102030405060708090a0b0c0d0e0f1011121314151617"
+    output_192 = "dda97ca4864cdfe06eaf70a0ec0d7191"
+    assert(Cipher(input_192, key_192) == output_192)
+
+    input_256 = "00112233445566778899aabbccddeeff"
+    key_256 = "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f"
+    output_256 = "8ea2b7ca516745bfeafc49904b496089"
+    assert(Cipher(input_256, key_256) == output_256)
+    
 if __name__ == "__main__":
     logging.basicConfig(
         level=logging.DEBUG,