1a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch# Authors:
2a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch#   Bram Cohen
3a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch#   Trevor Perrin - various changes
4a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch#
5a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch# See the LICENSE file for legal information regarding use of this file.
6a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch# Also see Bram Cohen's statement below
7a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch
85821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)"""
95821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)A pure python (slow) implementation of rijndael with a decent interface
105821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
115821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)To include -
125821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
135821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)from rijndael import rijndael
145821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
155821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)To do a key setup -
165821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
175821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)r = rijndael(key, block_size = 16)
185821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
195821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)key must be a string of length 16, 24, or 32
205821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)blocksize must be 16, 24, or 32. Default is 16
215821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
225821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)To use -
235821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
245821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)ciphertext = r.encrypt(plaintext)
255821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)plaintext = r.decrypt(ciphertext)
265821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
275821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)If any strings are of the wrong length a ValueError is thrown
285821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)"""
295821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
305821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# ported from the Java reference code by Bram Cohen, bram@gawth.com, April 2001
315821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# this code is public domain, unless someone makes
325821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# an intellectual property claim against the reference
335821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# code, in which case it can be made public domain by
345821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# deleting all the comments and renaming all the variables
355821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
365821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)import copy
375821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)import string
385821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
395821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)shifts = [[[0, 0], [1, 3], [2, 2], [3, 1]],
405821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)          [[0, 0], [1, 5], [2, 4], [3, 3]],
415821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)          [[0, 0], [1, 7], [3, 5], [4, 4]]]
425821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
435821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# [keysize][block_size]
445821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)num_rounds = {16: {16: 10, 24: 12, 32: 14}, 24: {16: 12, 24: 12, 32: 14}, 32: {16: 14, 24: 14, 32: 14}}
455821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
465821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)A = [[1, 1, 1, 1, 1, 0, 0, 0],
475821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)     [0, 1, 1, 1, 1, 1, 0, 0],
485821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)     [0, 0, 1, 1, 1, 1, 1, 0],
495821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)     [0, 0, 0, 1, 1, 1, 1, 1],
505821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)     [1, 0, 0, 0, 1, 1, 1, 1],
515821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)     [1, 1, 0, 0, 0, 1, 1, 1],
525821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)     [1, 1, 1, 0, 0, 0, 1, 1],
535821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)     [1, 1, 1, 1, 0, 0, 0, 1]]
545821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
555821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# produce log and alog tables, needed for multiplying in the
565821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# field GF(2^m) (generator = 3)
575821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)alog = [1]
58a02191e04bc25c4935f804f2c080ae28663d096dBen Murdochfor i in range(255):
595821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    j = (alog[-1] << 1) ^ alog[-1]
605821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    if j & 0x100 != 0:
615821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        j ^= 0x11B
625821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    alog.append(j)
635821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
645821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)log = [0] * 256
65a02191e04bc25c4935f804f2c080ae28663d096dBen Murdochfor i in range(1, 255):
665821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    log[alog[i]] = i
675821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
685821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# multiply two elements of GF(2^m)
695821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)def mul(a, b):
705821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    if a == 0 or b == 0:
715821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        return 0
725821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    return alog[(log[a & 0xFF] + log[b & 0xFF]) % 255]
735821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
745821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# substitution box based on F^{-1}(x)
75a02191e04bc25c4935f804f2c080ae28663d096dBen Murdochbox = [[0] * 8 for i in range(256)]
765821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)box[1][7] = 1
77a02191e04bc25c4935f804f2c080ae28663d096dBen Murdochfor i in range(2, 256):
785821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    j = alog[255 - log[i]]
79a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch    for t in range(8):
805821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        box[i][t] = (j >> (7 - t)) & 0x01
815821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
825821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)B = [0, 1, 1, 0, 0, 0, 1, 1]
835821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
845821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# affine transform:  box[i] <- B + A*box[i]
85a02191e04bc25c4935f804f2c080ae28663d096dBen Murdochcox = [[0] * 8 for i in range(256)]
86a02191e04bc25c4935f804f2c080ae28663d096dBen Murdochfor i in range(256):
87a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch    for t in range(8):
885821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        cox[i][t] = B[t]
89a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        for j in range(8):
905821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            cox[i][t] ^= A[t][j] * box[i][j]
915821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
925821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# S-boxes and inverse S-boxes
935821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)S =  [0] * 256
945821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)Si = [0] * 256
95a02191e04bc25c4935f804f2c080ae28663d096dBen Murdochfor i in range(256):
965821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    S[i] = cox[i][0] << 7
97a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch    for t in range(1, 8):
985821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        S[i] ^= cox[i][t] << (7-t)
995821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    Si[S[i] & 0xFF] = i
1005821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
1015821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# T-boxes
1025821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)G = [[2, 1, 1, 3],
1035821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    [3, 2, 1, 1],
1045821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    [1, 3, 2, 1],
1055821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    [1, 1, 3, 2]]
1065821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
107a02191e04bc25c4935f804f2c080ae28663d096dBen MurdochAA = [[0] * 8 for i in range(4)]
1085821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
109a02191e04bc25c4935f804f2c080ae28663d096dBen Murdochfor i in range(4):
110a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch    for j in range(4):
1115821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        AA[i][j] = G[i][j]
1125821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        AA[i][i+4] = 1
1135821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
114a02191e04bc25c4935f804f2c080ae28663d096dBen Murdochfor i in range(4):
1155821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    pivot = AA[i][i]
1165821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    if pivot == 0:
1175821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        t = i + 1
1185821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        while AA[t][i] == 0 and t < 4:
1195821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            t += 1
1205821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            assert t != 4, 'G matrix must be invertible'
121a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch            for j in range(8):
1225821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                AA[i][j], AA[t][j] = AA[t][j], AA[i][j]
1235821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            pivot = AA[i][i]
124a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch    for j in range(8):
1255821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        if AA[i][j] != 0:
1265821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            AA[i][j] = alog[(255 + log[AA[i][j] & 0xFF] - log[pivot & 0xFF]) % 255]
127a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch    for t in range(4):
1285821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        if i != t:
129a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch            for j in range(i+1, 8):
1305821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                AA[t][j] ^= mul(AA[i][j], AA[t][i])
1315821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            AA[t][i] = 0
1325821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
133a02191e04bc25c4935f804f2c080ae28663d096dBen MurdochiG = [[0] * 4 for i in range(4)]
1345821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
135a02191e04bc25c4935f804f2c080ae28663d096dBen Murdochfor i in range(4):
136a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch    for j in range(4):
1375821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        iG[i][j] = AA[i][j + 4]
1385821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
1395821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)def mul4(a, bs):
1405821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    if a == 0:
1415821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        return 0
1425821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    r = 0
1435821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    for b in bs:
1445821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        r <<= 8
1455821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        if b != 0:
1465821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            r = r | mul(a, b)
1475821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    return r
1485821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
1495821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)T1 = []
1505821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)T2 = []
1515821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)T3 = []
1525821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)T4 = []
1535821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)T5 = []
1545821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)T6 = []
1555821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)T7 = []
1565821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)T8 = []
1575821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)U1 = []
1585821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)U2 = []
1595821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)U3 = []
1605821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)U4 = []
1615821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
162a02191e04bc25c4935f804f2c080ae28663d096dBen Murdochfor t in range(256):
1635821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    s = S[t]
1645821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    T1.append(mul4(s, G[0]))
1655821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    T2.append(mul4(s, G[1]))
1665821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    T3.append(mul4(s, G[2]))
1675821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    T4.append(mul4(s, G[3]))
1685821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
1695821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    s = Si[t]
1705821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    T5.append(mul4(s, iG[0]))
1715821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    T6.append(mul4(s, iG[1]))
1725821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    T7.append(mul4(s, iG[2]))
1735821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    T8.append(mul4(s, iG[3]))
1745821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
1755821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    U1.append(mul4(t, iG[0]))
1765821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    U2.append(mul4(t, iG[1]))
1775821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    U3.append(mul4(t, iG[2]))
1785821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    U4.append(mul4(t, iG[3]))
1795821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
1805821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)# round constants
1815821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)rcon = [1]
1825821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)r = 1
183a02191e04bc25c4935f804f2c080ae28663d096dBen Murdochfor t in range(1, 30):
1845821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    r = mul(2, r)
1855821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    rcon.append(r)
1865821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
1875821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del A
1885821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del AA
1895821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del pivot
1905821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del B
1915821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del G
1925821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del box
1935821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del log
1945821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del alog
1955821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del i
1965821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del j
1975821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del r
1985821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del s
1995821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del t
2005821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del mul
2015821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del mul4
2025821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del cox
2035821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)del iG
2045821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
2055821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)class rijndael:
2065821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    def __init__(self, key, block_size = 16):
2075821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        if block_size != 16 and block_size != 24 and block_size != 32:
2085821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            raise ValueError('Invalid block size: ' + str(block_size))
2095821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        if len(key) != 16 and len(key) != 24 and len(key) != 32:
2105821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            raise ValueError('Invalid key size: ' + str(len(key)))
2115821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        self.block_size = block_size
2125821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
2135821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        ROUNDS = num_rounds[len(key)][block_size]
214a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        BC = block_size // 4
2155821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        # encryption round keys
216a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        Ke = [[0] * BC for i in range(ROUNDS + 1)]
2175821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        # decryption round keys
218a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        Kd = [[0] * BC for i in range(ROUNDS + 1)]
2195821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        ROUND_KEY_COUNT = (ROUNDS + 1) * BC
220a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        KC = len(key) // 4
2215821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
2225821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        # copy user material bytes into temporary ints
2235821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        tk = []
224a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        for i in range(0, KC):
225a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch            tk.append((key[i * 4] << 24) | (key[i * 4 + 1] << 16) |
226a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                (key[i * 4 + 2] << 8) | key[i * 4 + 3])
2275821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
2285821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        # copy values into round key arrays
2295821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        t = 0
2305821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        j = 0
2315821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        while j < KC and t < ROUND_KEY_COUNT:
232a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch            Ke[t // BC][t % BC] = tk[j]
233a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch            Kd[ROUNDS - (t // BC)][t % BC] = tk[j]
2345821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            j += 1
2355821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            t += 1
2365821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        tt = 0
2375821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        rconpointer = 0
2385821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        while t < ROUND_KEY_COUNT:
2395821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            # extrapolate using phi (the round key evolution function)
2405821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            tt = tk[KC - 1]
2415821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            tk[0] ^= (S[(tt >> 16) & 0xFF] & 0xFF) << 24 ^  \
2425821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                     (S[(tt >>  8) & 0xFF] & 0xFF) << 16 ^  \
2435821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                     (S[ tt        & 0xFF] & 0xFF) <<  8 ^  \
2445821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                     (S[(tt >> 24) & 0xFF] & 0xFF)       ^  \
2455821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                     (rcon[rconpointer]    & 0xFF) << 24
2465821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            rconpointer += 1
2475821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            if KC != 8:
248a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                for i in range(1, KC):
2495821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                    tk[i] ^= tk[i-1]
2505821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            else:
251a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                for i in range(1, KC // 2):
2525821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                    tk[i] ^= tk[i-1]
253a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                tt = tk[KC // 2 - 1]
254a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                tk[KC // 2] ^= (S[ tt        & 0xFF] & 0xFF)       ^ \
2555821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                              (S[(tt >>  8) & 0xFF] & 0xFF) <<  8 ^ \
2565821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                              (S[(tt >> 16) & 0xFF] & 0xFF) << 16 ^ \
2575821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                              (S[(tt >> 24) & 0xFF] & 0xFF) << 24
258a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                for i in range(KC // 2 + 1, KC):
2595821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                    tk[i] ^= tk[i-1]
2605821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            # copy values into round key arrays
2615821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            j = 0
2625821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            while j < KC and t < ROUND_KEY_COUNT:
263a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                Ke[t // BC][t % BC] = tk[j]
264a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                Kd[ROUNDS - (t // BC)][t % BC] = tk[j]
2655821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                j += 1
2665821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                t += 1
2675821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        # inverse MixColumn where needed
268a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        for r in range(1, ROUNDS):
269a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch            for j in range(BC):
2705821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                tt = Kd[r][j]
2715821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                Kd[r][j] = U1[(tt >> 24) & 0xFF] ^ \
2725821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                           U2[(tt >> 16) & 0xFF] ^ \
2735821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                           U3[(tt >>  8) & 0xFF] ^ \
2745821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                           U4[ tt        & 0xFF]
2755821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        self.Ke = Ke
2765821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        self.Kd = Kd
2775821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
2785821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    def encrypt(self, plaintext):
2795821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        if len(plaintext) != self.block_size:
2805821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            raise ValueError('wrong block length, expected ' + str(self.block_size) + ' got ' + str(len(plaintext)))
2815821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        Ke = self.Ke
2825821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
283a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        BC = self.block_size // 4
2845821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        ROUNDS = len(Ke) - 1
2855821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        if BC == 4:
2865821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            SC = 0
2875821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        elif BC == 6:
2885821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            SC = 1
2895821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        else:
2905821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            SC = 2
2915821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        s1 = shifts[SC][1][0]
2925821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        s2 = shifts[SC][2][0]
2935821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        s3 = shifts[SC][3][0]
2945821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        a = [0] * BC
2955821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        # temporary work array
2965821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        t = []
2975821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        # plaintext to ints + key
298a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        for i in range(BC):
299a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch            t.append((plaintext[i * 4    ] << 24 |
300a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                      plaintext[i * 4 + 1] << 16 |
301a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                      plaintext[i * 4 + 2] <<  8 |
302a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                      plaintext[i * 4 + 3]        ) ^ Ke[0][i])
3035821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        # apply round transforms
304a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        for r in range(1, ROUNDS):
305a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch            for i in range(BC):
3065821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                a[i] = (T1[(t[ i           ] >> 24) & 0xFF] ^
3075821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                        T2[(t[(i + s1) % BC] >> 16) & 0xFF] ^
3085821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                        T3[(t[(i + s2) % BC] >>  8) & 0xFF] ^
3095821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                        T4[ t[(i + s3) % BC]        & 0xFF]  ) ^ Ke[r][i]
3105821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            t = copy.copy(a)
3115821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        # last round is special
3125821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        result = []
313a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        for i in range(BC):
3145821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            tt = Ke[ROUNDS][i]
3155821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            result.append((S[(t[ i           ] >> 24) & 0xFF] ^ (tt >> 24)) & 0xFF)
3165821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            result.append((S[(t[(i + s1) % BC] >> 16) & 0xFF] ^ (tt >> 16)) & 0xFF)
3175821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            result.append((S[(t[(i + s2) % BC] >>  8) & 0xFF] ^ (tt >>  8)) & 0xFF)
3185821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            result.append((S[ t[(i + s3) % BC]        & 0xFF] ^  tt       ) & 0xFF)
319a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        return bytearray(result)
3205821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
3215821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    def decrypt(self, ciphertext):
3225821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        if len(ciphertext) != self.block_size:
3235821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            raise ValueError('wrong block length, expected ' + str(self.block_size) + ' got ' + str(len(plaintext)))
3245821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        Kd = self.Kd
3255821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
326a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        BC = self.block_size // 4
3275821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        ROUNDS = len(Kd) - 1
3285821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        if BC == 4:
3295821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            SC = 0
3305821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        elif BC == 6:
3315821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            SC = 1
3325821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        else:
3335821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            SC = 2
3345821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        s1 = shifts[SC][1][1]
3355821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        s2 = shifts[SC][2][1]
3365821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        s3 = shifts[SC][3][1]
3375821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        a = [0] * BC
3385821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        # temporary work array
3395821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        t = [0] * BC
3405821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        # ciphertext to ints + key
341a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        for i in range(BC):
342a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch            t[i] = (ciphertext[i * 4    ] << 24 |
343a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                    ciphertext[i * 4 + 1] << 16 |
344a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                    ciphertext[i * 4 + 2] <<  8 |
345a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch                    ciphertext[i * 4 + 3]        ) ^ Kd[0][i]
3465821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        # apply round transforms
347a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        for r in range(1, ROUNDS):
348a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch            for i in range(BC):
3495821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                a[i] = (T5[(t[ i           ] >> 24) & 0xFF] ^
3505821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                        T6[(t[(i + s1) % BC] >> 16) & 0xFF] ^
3515821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                        T7[(t[(i + s2) % BC] >>  8) & 0xFF] ^
3525821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)                        T8[ t[(i + s3) % BC]        & 0xFF]  ) ^ Kd[r][i]
3535821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            t = copy.copy(a)
3545821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        # last round is special
3555821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        result = []
356a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        for i in range(BC):
3575821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            tt = Kd[ROUNDS][i]
3585821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            result.append((Si[(t[ i           ] >> 24) & 0xFF] ^ (tt >> 24)) & 0xFF)
3595821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            result.append((Si[(t[(i + s1) % BC] >> 16) & 0xFF] ^ (tt >> 16)) & 0xFF)
3605821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            result.append((Si[(t[(i + s2) % BC] >>  8) & 0xFF] ^ (tt >>  8)) & 0xFF)
3615821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)            result.append((Si[ t[(i + s3) % BC]        & 0xFF] ^  tt       ) & 0xFF)
362a02191e04bc25c4935f804f2c080ae28663d096dBen Murdoch        return bytearray(result)
3635821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
3645821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)def encrypt(key, block):
3655821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    return rijndael(key, len(block)).encrypt(block)
3665821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
3675821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)def decrypt(key, block):
3685821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    return rijndael(key, len(block)).decrypt(block)
3695821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
3705821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)def test():
3715821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    def t(kl, bl):
3725821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        b = 'b' * bl
3735821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        r = rijndael('a' * kl, bl)
3745821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)        assert r.decrypt(r.encrypt(b)) == b
3755821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    t(16, 16)
3765821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    t(16, 24)
3775821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    t(16, 32)
3785821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    t(24, 16)
3795821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    t(24, 24)
3805821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    t(24, 32)
3815821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    t(32, 16)
3825821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    t(32, 24)
3835821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)    t(32, 32)
3845821806d5e7f356e8fa4b058a389a808ea183019Torne (Richard Coles)
385