1import unittest
2from util import *
3
4cases = [
5    # AES test vectors from FIPS 197.
6    [ 128, "000102030405060708090a0b0c0d0e0f",
7           "00112233445566778899aabbccddeeff",
8           "69c4e0d86a7b0430d8cdb78070b4c55a" ],
9    [ 192, "000102030405060708090a0b0c0d0e0f1011121314151617",
10           "00112233445566778899aabbccddeeff",
11           "dda97ca4864cdfe06eaf70a0ec0d7191" ],
12    [ 256, "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f",
13           "00112233445566778899aabbccddeeff",
14           "8ea2b7ca516745bfeafc49904b496089" ],
15    # AES-ECB test vectors from NIST sp800-38a.
16    [ 128, "2b7e151628aed2a6abf7158809cf4f3c",
17           "6bc1bee22e409f96e93d7e117393172a",
18           "3ad77bb40d7a3660a89ecaf32466ef97" ],
19    [ 128, "2b7e151628aed2a6abf7158809cf4f3c",
20           "ae2d8a571e03ac9c9eb76fac45af8e51",
21           "f5d3d58503b9699de785895a96fdbaaf" ],
22    [ 128, "2b7e151628aed2a6abf7158809cf4f3c",
23           "30c81c46a35ce411e5fbc1191a0a52ef",
24           "43b1cd7f598ece23881b00e3ed030688" ],
25    [ 128, "2b7e151628aed2a6abf7158809cf4f3c",
26           "f69f2445df4f9b17ad2b417be66c3710",
27           "7b0c785e27e8ad3f8223207104725dd4" ],
28    [ 192, "8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b",
29           "6bc1bee22e409f96e93d7e117393172a",
30           "bd334f1d6e45f25ff712a214571fa5cc" ],
31    [ 192, "8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b",
32           "ae2d8a571e03ac9c9eb76fac45af8e51",
33           "974104846d0ad3ad7734ecb3ecee4eef" ],
34    [ 192, "8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b",
35           "30c81c46a35ce411e5fbc1191a0a52ef",
36           "ef7afd2270e2e60adce0ba2face6444e" ],
37    [ 192, "8e73b0f7da0e6452c810f32b809079e562f8ead2522c6b7b",
38           "f69f2445df4f9b17ad2b417be66c3710",
39           "9a4b41ba738d6c72fb16691603c18e0e" ],
40    [ 256, "603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4",
41           "6bc1bee22e409f96e93d7e117393172a",
42           "f3eed1bdb5d2a03c064b5a7e3db181f8" ],
43    [ 256, "603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4",
44           "ae2d8a571e03ac9c9eb76fac45af8e51",
45           "591ccb10d410ed26dc5ba74a31362870" ],
46    [ 256, "603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4",
47           "30c81c46a35ce411e5fbc1191a0a52ef",
48           "b6ed21b99ca6f4f9f153e7b1beafed1d" ],
49    [ 256, "603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4",
50           "f69f2445df4f9b17ad2b417be66c3710",
51           "23304b7a39f9f3ff067d8d8f9e24ecc7" ],
52]
53
54class AESTests(unittest.TestCase):
55
56    ENCRYPT, DECRYPT = 1, 2
57
58    def test_aes(self):
59
60        for c in cases:
61            key, plain, cypher = [make_cbuffer(s)[0] for s in c[1:]]
62            key_bytes = { 128: 16, 192: 24, 256: 32}[c[0]]
63            self.assertEqual(len(key), key_bytes)
64
65            for p, f, o in [(plain,  self.ENCRYPT, cypher),
66                            (cypher, self.DECRYPT, plain)]:
67
68                out_buf, out_len = make_cbuffer('00' * len(o))
69                ret = wally_aes(key, len(key), p, len(p), f, out_buf, out_len)
70                self.assertEqual(ret, 0)
71                self.assertEqual(h(out_buf), h(o))
72
73
74    def get_cbc_cases(self):
75        lines = []
76        with open(root_dir + 'src/data/aes-cbc-pkcs7.txt', 'r') as f:
77            for l in f.readlines():
78                if len(l.strip()) and not l.startswith('#'):
79                    lines.append(l.strip().split('=')[1])
80        return [lines[x:x+4] for x in range(0, len(lines), 4)]
81
82    def test_aes_cbc(self):
83        for c in self.get_cbc_cases():
84            plain, key, iv, cypher = [make_cbuffer(s)[0] for s in c]
85
86            for p, f, o in [(plain,  self.ENCRYPT, cypher),
87                            (cypher, self.DECRYPT, plain)]:
88
89                out_buf, out_len = make_cbuffer('00' * len(o))
90                ret, written = wally_aes_cbc(key, len(key), iv, len(iv),
91                                             p, len(p), f, out_buf, out_len)
92                self.assertEqual((ret, written), (0, len(o)))
93                self.assertEqual(h(out_buf), h(o))
94
95
96if __name__ == '__main__':
97    unittest.main()
98