diff --git a/aes.py b/aes.py index c569372..5f26aed 100755 --- a/aes.py +++ b/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]) diff --git a/fuzz_test.py b/fuzz_test.py new file mode 100644 index 0000000..bd1a10f --- /dev/null +++ b/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) diff --git a/test.py b/test.py index 42aaa4a..432cc36 100644 --- a/test.py +++ b/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() \ No newline at end of file