[U-Boot] [PATCH V2 4/4] pytest: aes: add test for aes192 and aes256

[Philippe Reynes]

This commit update the aes tests to check the
aes192 and aes256.

Signed-off-by: Philippe Reynes <philippe.reynes at softathome.com>
---
 test/py/tests/test_aes.py | 118 +++++++++++++++++++++++++++++++++++-----------
 1 file changed, 91 insertions(+), 27 deletions(-)

Changelog:
v2:
- add a really simple commit text
- re-write the test with the new version of test for aes128
- add a test to check that results of aes128/192/256 are differents

diff --git a/test/py/tests/test_aes.py b/test/py/tests/test_aes.py
index 01d514e..8f4cad3 100644
--- a/test/py/tests/test_aes.py
+++ b/test/py/tests/test_aes.py
@@ -11,82 +11,132 @@ def test_aes(u_boot_console):
     cipher and uncipher a random data
     """
 
-    def generate_key_iv_data():
+    class Seed:
+        def __init__(self):
+            self.seed = 0
+
+        def get(self):
+            self.seed = self.seed + 1
+            return self.seed
+
+    def generate_random_data(addr, size):
+        output = cons.run_command('random %d 0x%x %x' % (addr, size, seed.get()))
+        assert('%d bytes filled with random data' % size in ''.join(output))
+
+    def generate_key_iv_data(key_size):
         # Generate random key
-        output = cons.run_command('random 1000 0x10 0')
-        assert('16 bytes filled with random data' in ''.join(output))
+        generate_random_data(1000, key_size)
 
         # Generate random IV
-        output = cons.run_command('random 2000 0x10 0')
-        assert('16 bytes filled with random data' in ''.join(output))
+        generate_random_data(2000, 16)
 
         # Generate random data
-        output = cons.run_command('random 3000 0x20 0')
-        assert('32 bytes filled with random data' in ''.join(output))
+        generate_random_data(3000, 32)
+
+    def test_nominal(aes_size):
+        key_size = aes_size / 8
 
-    def test_nominal():
         # Generate random key, iv and data
-        generate_key_iv_data()
+        generate_key_iv_data(key_size)
 
         # Encrypt random data
-        output = cons.run_command('aes enc 1000 2000 3000 4000 0x20')
+        output = cons.run_command('aes.%d enc 1000 2000 3000 4000 0x20' % aes_size)
 
         # Check that ciphered data are different than unciphered data
         output = cons.run_command('cmp.b 3000 4000 0x20')
         assert('Total of 0 byte(s) were the same' in ''.join(output))
 
         # Decrypt ciphered data
-        output = cons.run_command('aes dec 1000 2000 4000 5000 0x20')
+        output = cons.run_command('aes.%d dec 1000 2000 4000 5000 0x20' % aes_size)
 
         # Check that unciphered data are the same than initial data
         output = cons.run_command('cmp.b 3000 5000 0x20')
         assert('Total of 32 byte(s) were the same' in ''.join(output))
 
-    def test_corrupted_key(seed):
+    def test_corrupted_key(aes_size):
+        key_size = aes_size / 8
+
         # Generate random key, iv and data
-        generate_key_iv_data()
+        generate_key_iv_data(key_size)
 
         # Encrypt random data
-        output = cons.run_command('aes enc 1000 2000 3000 4000 0x20')
+        output = cons.run_command('aes.%d enc 1000 2000 3000 4000 0x20' % aes_size)
 
         # Check that ciphered data are different than unciphered data
         output = cons.run_command('cmp.b 3000 4000 0x20')
         assert('Total of 0 byte(s) were the same' in ''.join(output))
 
         # Corrupt the key (simply generate a new one)
-        output = cons.run_command('random 1000 0x10 %d' % seed)
-        assert('16 bytes filled with random data' in ''.join(output))
+        generate_random_data(1000, key_size)
 
         # Decrypt ciphered data
-        output = cons.run_command('aes dec 1000 2000 4000 5000 0x20')
+        output = cons.run_command('aes.%d dec 1000 2000 4000 5000 0x20' % aes_size)
 
         # Check that unciphered data are different than initial data
         output = cons.run_command('cmp.b 3000 5000 0x20')
         assert('Total of 32 byte(s) were the same' not in ''.join(output))
 
-    def test_corrupted_iv(seed):
+    def test_corrupted_iv(aes_size):
         # Generate random key, iv and data
-        generate_key_iv_data()
+        generate_key_iv_data(aes_size / 8)
 
         # Encrypt random data
-        output = cons.run_command('aes enc 1000 2000 3000 4000 0x20')
+        output = cons.run_command('aes.%d enc 1000 2000 3000 4000 0x20' % aes_size)
 
         # Check that ciphered data are different than unciphered data
         output = cons.run_command('cmp.b 3000 4000 0x20')
         assert('Total of 0 byte(s) were the same' in ''.join(output))
 
         # Corrupt the iv (simply generate a new one)
-        output = cons.run_command('random 2000 0x10 %d' % seed)
-        assert('16 bytes filled with random data' in ''.join(output))
+        generate_random_data(2000, 16)
 
         # Decrypt ciphered data
-        output = cons.run_command('aes dec 1000 2000 4000 5000 0x20')
+        output = cons.run_command('aes.%d dec 1000 2000 4000 5000 0x20' % aes_size)
 
         # Check that unciphered data are different than initial data
         output = cons.run_command('cmp.b 3000 5000 0x20')
         assert('Total of 32 byte(s) were the same' not in ''.join(output))
 
+    def test_different_size():
+        # Generate random key, iv and data
+        generate_key_iv_data(256 / 8)
+
+        # Encrypt random data with aes128
+        output = cons.run_command('aes.%d enc 1000 2000 3000 4000 0x20' % 128)
+
+        # Check that ciphered data are different than unciphered data
+        output = cons.run_command('cmp.b 3000 4000 0x20')
+        assert('Total of 0 byte(s) were the same' in ''.join(output))
+
+        # Encrypt random data with aes192
+        output = cons.run_command('aes.%d enc 1000 2000 3000 5000 0x20' % 192)
+
+        # Check that ciphered data are different than unciphered data
+        output = cons.run_command('cmp.b 3000 5000 0x20')
+        assert('Total of 0 byte(s) were the same' in ''.join(output))
+
+        # Encrypt random data with aes256
+        output = cons.run_command('aes.%d enc 1000 2000 3000 6000 0x20' % 256)
+
+        # Check that ciphered data are different than initial data
+        output = cons.run_command('cmp.b 3000 6000 0x20')
+        assert('Total of 0 byte(s) were the same' in ''.join(output))
+
+        # Check that result of aes128 is different thant aes192
+        output = cons.run_command('cmp.b 4000 5000 0x20')
+        assert('Total of 32 byte(s) were the same' not in ''.join(output))
+
+        # Check that result of aes128 is different thant aes256
+        output = cons.run_command('cmp.b 4000 6000 0x20')
+        assert('Total of 32 byte(s) were the same' not in ''.join(output))
+
+        # Check that result of aes192 is different thant aes256
+        output = cons.run_command('cmp.b 5000 6000 0x20')
+        assert('Total of 32 byte(s) were the same' not in ''.join(output))
+
+
     cons = u_boot_console
+    seed = Seed()
 
     # Check that the option cmd_aes is enabled in the config
     if cons.config.buildconfig.get('config_cmd_aes', 'n') != 'y':
@@ -94,8 +144,22 @@ def test_aes(u_boot_console):
 
     # Send a command with no argument ...
     output = cons.run_command('aes')
-    assert('AES 128 CBC encryption' in ''.join(output))
+    assert('AES 128/192/256 CBC encryption' in ''.join(output))
+
+    # test aes128
+    test_nominal(128)
+    test_corrupted_key(128)
+    test_corrupted_iv(128)
+
+    # test aes192
+    test_nominal(192)
+    test_corrupted_key(192)
+    test_corrupted_iv(192)
+
+    # test aes256
+    test_nominal(256)
+    test_corrupted_key(256)
+    test_corrupted_iv(256)
 
-    test_nominal()
-    test_corrupted_key(666)
-    test_corrupted_iv(666)
+    # test that each key size provide different result
+    test_different_size()
-- 
2.7.4