from unittest import TestCase, main
import aes

def bytes(hexbytes):
    return [int(hexbyte, 16) for hexbyte in hexbytes.split()]

def chars(hexbytes):
    return ''.join(map(chr, bytes(hexbytes)))

class TestAES(TestCase):
    def test_KeyExpansion(self):
        equal = self.assertEqual

        # Appendix A.1 of FIPS 197, page 27.
        K = 0x2b7e151628aed2a6abf7158809cf4f3c
        equal(aes.KeyExpansion(K, 4, 4, 10),
              [0x2b7e1516, 0x28aed2a6, 0xabf71588, 0x09cf4f3c,
               0xa0fafe17, 0x88542cb1, 0x23a33939, 0x2a6c7605,
               0xf2c295f2, 0x7a96b943, 0x5935807a, 0x7359f67f,
               0x3d80477d, 0x4716fe3e, 0x1e237e44, 0x6d7a883b,
               0xef44a541, 0xa8525b7f, 0xb671253b, 0xdb0bad00,
               0xd4d1c6f8, 0x7c839d87, 0xcaf2b8bc, 0x11f915bc,
               0x6d88a37a, 0x110b3efd, 0xdbf98641, 0xca0093fd,
               0x4e54f70e, 0x5f5fc9f3, 0x84a64fb2, 0x4ea6dc4f,
               0xead27321, 0xb58dbad2, 0x312bf560, 0x7f8d292f,
               0xac7766f3, 0x19fadc21, 0x28d12941, 0x575c006e,
               0xd014f9a8, 0xc9ee2589, 0xe13f0cc8, 0xb6630ca6])

        # Appendix A.2 of FIPS 197, page 28.
        K = 0x8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b
        equal(aes.KeyExpansion(K, 6, 4, 12),
              [0x8e73b0f7, 0xda0e6452, 0xc810f32b, 0x809079e5,
               0x62f8ead2, 0x522c6b7b, 0xfe0c91f7, 0x2402f5a5,
               0xec12068e, 0x6c827f6b, 0x0e7a95b9, 0x5c56fec2,
               0x4db7b4bd, 0x69b54118, 0x85a74796, 0xe92538fd,
               0xe75fad44, 0xbb095386, 0x485af057, 0x21efb14f,
               0xa448f6d9, 0x4d6dce24, 0xaa326360, 0x113b30e6,
               0xa25e7ed5, 0x83b1cf9a, 0x27f93943, 0x6a94f767,
               0xc0a69407, 0xd19da4e1, 0xec1786eb, 0x6fa64971,
               0x485f7032, 0x22cb8755, 0xe26d1352, 0x33f0b7b3,
               0x40beeb28, 0x2f18a259, 0x6747d26b, 0x458c553e,
               0xa7e1466c, 0x9411f1df, 0x821f750a, 0xad07d753,
               0xca400538, 0x8fcc5006, 0x282d166a, 0xbc3ce7b5,
               0xe98ba06f, 0x448c773c, 0x8ecc7204, 0x01002202])

        # Appendix A.3 of FIPS 197, page 30.
        K = 0x603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4
        equal(aes.KeyExpansion(K, 8, 4, 14),
              [0x603deb10, 0x15ca71be, 0x2b73aef0, 0x857d7781,
               0x1f352c07, 0x3b6108d7, 0x2d9810a3, 0x0914dff4,
               0x9ba35411, 0x8e6925af, 0xa51a8b5f, 0x2067fcde,
               0xa8b09c1a, 0x93d194cd, 0xbe49846e, 0xb75d5b9a,
               0xd59aecb8, 0x5bf3c917, 0xfee94248, 0xde8ebe96,
               0xb5a9328a, 0x2678a647, 0x98312229, 0x2f6c79b3,
               0x812c81ad, 0xdadf48ba, 0x24360af2, 0xfab8b464,
               0x98c5bfc9, 0xbebd198e, 0x268c3ba7, 0x09e04214,
               0x68007bac, 0xb2df3316, 0x96e939e4, 0x6c518d80,
               0xc814e204, 0x76a9fb8a, 0x5025c02d, 0x59c58239,
               0xde136967, 0x6ccc5a71, 0xfa256395, 0x9674ee15,
               0x5886ca5d, 0x2e2f31d7, 0x7e0af1fa, 0x27cf73c3,
               0x749c47ab, 0x18501dda, 0xe2757e4f, 0x7401905a,
               0xcafaaae3, 0xe4d59b34, 0x9adf6ace, 0xbd10190d,
               0xfe4890d1, 0xe6188d0b, 0x046df344, 0x706c631e])

    def test_cipher(self):
        equal = self.assertEqual

        # Appendix B of FIPS 197, page 33.
        K = 0x2b7e151628aed2a6abf7158809cf4f3c
        plaintext = bytes('32 43 f6 a8 88 5a 30 8d 31 31 98 a2 e0 37 07 34')
        ciphertext = bytes('39 25 84 1d 02 dc 09 fb dc 11 85 97 19 6a 0b 32')
        equal(aes.AES(128, K).encipher(plaintext), ciphertext)
        equal(aes.AES(128, K).decipher(ciphertext), plaintext)

        # Appendix C of FIPS 197, page 35.
        K = 0x000102030405060708090a0b0c0d0e0f
        plaintext = bytes('00 11 22 33 44 55 66 77 88 99 aa bb cc dd ee ff')
        ciphertext = bytes('69 c4 e0 d8 6a 7b 04 30 d8 cd b7 80 70 b4 c5 5a')
        equal(aes.AES(128, K).encipher(plaintext), ciphertext)
        equal(aes.AES(128, K).decipher(ciphertext), plaintext)

        K = 0x000102030405060708090a0b0c0d0e0f1011121314151617
        plaintext = bytes('00 11 22 33 44 55 66 77 88 99 aa bb cc dd ee ff')
        ciphertext = bytes('dd a9 7c a4 86 4c df e0 6e af 70 a0 ec 0d 71 91')
        equal(aes.AES(192, K).encipher(plaintext), ciphertext)
        equal(aes.AES(192, K).decipher(ciphertext), plaintext)

        K = 0x000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f
        plaintext = bytes('00 11 22 33 44 55 66 77 88 99 aa bb cc dd ee ff')
        ciphertext = bytes('8e a2 b7 ca 51 67 45 bf ea fc 49 90 4b 49 60 89')
        equal(aes.AES(256, K).encipher(plaintext), ciphertext)
        equal(aes.AES(256, K).decipher(ciphertext), plaintext)

    def test_mac(self):
        equal = self.assertEqual

        # http://www.ietf.org/internet-drafts/draft-songlee-aes-cmac-03.txt
        K = 0x2b7e151628aed2a6abf7158809cf4f3c

        text = ''
        mac = 0xbb1d6929e95937287fa37d129b756746
        equal(aes.mac(aes.AES(128, K), text), mac)

        text = chars('6b c1 be e2 2e 40 9f 96 e9 3d 7e 11 73 93 17 2a')
        mac = 0x070a16b46b4d4144f79bdd9dd04a287c
        equal(aes.mac(aes.AES(128, K), text), mac)

        text = chars('''6b c1 be e2 2e 40 9f 96 e9 3d 7e 11 73 93 17 2a
                        ae 2d 8a 57 1e 03 ac 9c 9e b7 6f ac 45 af 8e 51
                        30 c8 1c 46 a3 5c e4 11''')
        mac = 0xdfa66747de9ae63030ca32611497c827
        equal(aes.mac(aes.AES(128, K), text), mac)

        text = chars('''6b c1 be e2 2e 40 9f 96 e9 3d 7e 11 73 93 17 2a  
                        ae 2d 8a 57 1e 03 ac 9c 9e b7 6f ac 45 af 8e 51  
                        30 c8 1c 46 a3 5c e4 11 e5 fb c1 19 1a 0a 52 ef  
                        f6 9f 24 45 df 4f 9b 17 ad 2b 41 7b e6 6c 37 10''')
        mac = 0x51f0bebf7e3b9d92fc49741779363cfe

if __name__ == '__main__':
    main()