aestools_tests.py

import unittest
import aestools


class AesToolsTestCase(unittest.TestCase):

    def test_padding(self):
        """ Adds specified padding to bytes """
        input = b'YELLOW SUBMARINE'
        desired_length = 20

        self.assertEqual(aestools.padded(input, desired_length), b'YELLOW SUBMARINE\x04\x04\x04\x04')

    def test_random_key(self):
        """ Creates random X byte aes key """
        key = aestools.random_key(16)
        keytwo = aestools.random_key(16)
        keythree = aestools.random_key(5)

        self.assertEqual(len(key), 16)
        self.assertEqual(len(keythree), 5)
        self.assertNotEqual(key, keytwo)

    def test_add_bytes(self):
        input = b'test'
        output = aestools.add_bytes_to_input(input)
        minlength = len(input) + 10
        maxlength = len(input) + 20

        self.assertNotEqual(input, output)
        self.assertTrue(minlength <= len(output) <= maxlength)

    def test_get_prefix(self):
        function = aestools.get_black_box(True)
        for prefix_length in range(0, 128):
            aestools.PREFIX = aestools.random_key(prefix_length)
            blocksize, length, padding = aestools.get_block_size_and_length(function)
            self.assertEqual(aestools.get_prefix_length(function, blocksize, length, padding), prefix_length)

    def test_quote_symbols(self):
        encrypted = aestools.cbc_black_box("test;admin=true;")
        self.assertFalse(aestools.has_admin_string(encrypted))

if __name__ == '__main__':
    unittest.main()