Advanced testing system + small changes in aes.py

aes.py

@@ -81,7 +81,7 @@ def text2matrix(text):
         if i % 4 == 0:
             matrix.append([byte])
         else:
-            matrix[i / 4].append(byte)
+            matrix[i // 4].append(byte)
     return matrix
 
 
@@ -95,6 +95,8 @@ def matrix2text(matrix):
 
 class AES:
     def __init__(self, master_key):
+        if master_key < 0:
+            raise Exception('Key can not be < 0')
         self.change_key(master_key)
 
     def change_key(self, master_key):
@@ -106,7 +108,7 @@ def change_key(self, master_key):
             if i % 4 == 0:
                 byte = self.round_keys[i - 4][0]        \
                      ^ Sbox[self.round_keys[i - 1][1]]  \
-                     ^ Rcon[i / 4]
+                     ^ Rcon[i // 4]
                 self.round_keys[i].append(byte)
 
                 for j in range(1, 4):
@@ -122,6 +124,8 @@ def change_key(self, master_key):
         # print self.round_keys
 
     def encrypt(self, plaintext):
+        if plaintext < 0:
+            raise Exception('Plaintext can not be < 0')
         self.plain_state = text2matrix(plaintext)
 
         self.__add_round_key(self.plain_state, self.round_keys[:4])

fuzz_test.py

@@ -0,0 +1,75 @@
+__author__ = 'L3odr0id'
+"""
+Fuzz testing
+"""
+import time
+from fuzzingbook.Fuzzer import RandomFuzzer
+from aes import AES
+
+
+def test_system(key, data):
+    """
+    Encrypt and decrypt @key and @data
+    """
+    a = AES(key)
+
+    plaintext = data
+    encrypted = a.encrypt(plaintext)
+    decrypted = a.decrypt(encrypted)
+
+    # check if data is encrypted
+    assert encrypted != decrypted, 'Data is not encrypted!'
+
+    return decrypted
+
+
+def fuzz_testing(num_of_tests=10):
+    """
+    Do tests
+    """
+    ok = True
+    start = time.clock()
+    errors = []
+
+    # generate integers. Max possible length is 38
+    r = RandomFuzzer(min_length=1, char_start=48, char_range=9, max_length=38)
+
+    dot = num_of_tests // 10    # decoration stuff
+
+    for i in range(0, num_of_tests):
+        if i % dot == 0:
+            print('.', end='')
+        # generate test case
+        key = int(r.fuzz())
+        data = int(r.fuzz())
+
+        # analyze output
+        try:
+            result = test_system(key, data)
+            if data != result:
+                errors.append('TEST ' + str(i) + ' FAILED\nKey = "' + str(key) + '"\nData = "' + str(data) + '"')
+                errors.append('Result = "'+str(result)+'"')
+                errors.append('Reason: Error in encryption - decryption process\n')
+                ok = False
+        except AssertionError:
+            errors.append('TEST '+str(i)+' FAILED\nKey = "'+str(key)+'"\nData = "'+str(data)+'"')
+            errors.append('Reason: Data was not encrypted\n')
+            ok = False
+
+    # print testing results
+    elapsed = time.clock()
+    elapsed = elapsed - start
+    print()
+    for i in errors:
+        print(i)
+    print('Ran ' + str(num_of_tests) + ' tests in '+str(round(elapsed, 3))+'s')
+    print()
+    if ok:
+        print('OK')
+    else:
+        print('FAILED')
+
+
+if __name__ == "__main__":
+    print('Fuzz test started.')
+    fuzz_testing(1337)

test.py

@@ -3,6 +3,7 @@
 import unittest
 from aes import AES
 
+
 class AES_TEST(unittest.TestCase):
     def setUp(self):
         master_key = 0x2b7e151628aed2a6abf7158809cf4f3c
@@ -20,5 +21,61 @@ def test_decryption(self):
 
         self.assertEqual(decrypted, 0x3243f6a8885a308d313198a2e0370734)
 
-if __name__ == '__main__':
+
+"""
+L3odr0id testing system
+
+https://github.com/L3odr0id
+"""
+from hypothesis import given, example
+from hypothesis.strategies import integers, text
+from aes import text2matrix, matrix2text
+
+
+class TestMatrix(unittest.TestCase):
+    @given(a=integers(min_value=0))
+    def test_matrix(self, a):
+        h = int(hex(a), 16)
+        text = text2matrix(h)
+        mat = matrix2text(text)
+        self.assertEqual(h, mat, '"There and back" rule does not work with matrix')
+
+    def test_incorrect_input(self):
+        with self.assertRaises(IndexError, msg='You should handle broken matrix input'):
+            matrix2text([[], []])
+
+
+class TestMainClass(unittest.TestCase):
+    @given(key=integers(min_value=0), data=integers(min_value=0))
+    def test_system(self, key, data):
+        a = AES(int(hex(key), 16))
+
+        plaintext = data
+        encrypted = a.encrypt(plaintext)
+        decrypted = a.decrypt(encrypted)
+
+        self.assertFalse(encrypted == decrypted, 'Data is not encrypted')
+        self.assertEqual(decrypted, data, 'Something is wrong with encryption - decryption process')
+
+    @given(key=text(), data=text())
+    def test_wrong_input(self, key, data):
+        with self.assertRaises(TypeError, msg='You should handle wrong input'):
+            a = AES(key)
+
+            plaintext = data
+            encrypted = a.encrypt(plaintext)
+            a.decrypt(encrypted)
+
+    @given(key=integers(max_value=-1))
+    def test_bad_input(self, key):
+        with self.assertRaises(Exception) as context:
+            AES(key)
+        self.assertTrue('Key can not be < 0' in str(context.exception))
+        a = AES(123)
+        with self.assertRaises(Exception) as context:
+            a.encrypt(key)
+        self.assertTrue('Plaintext can not be < 0' in str(context.exception))
+
+
+if __name__ == "__main__":
     unittest.main()