CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

1
import copy
2
import struct
3
import sys
4

5

6
def chunks(lst, n):
7
    for i in range(0, len(lst), n):
8
        yield lst[i:i + n]
9

10

11
def _cw(word):
12
    return (word[0] << 24) | (word[1] << 16) | (word[2] << 8) | word[3]
13

14

15
def _s2b(text):
16
    return list(ord(c)for c in text)
17

18

19
def _b2s(binary):
20
    return "".join(chr(b)for b in binary)
21

22

23
if sys.version_info[0] >= 3:
24
    xrange = range
25

26
    def _s2b(text):
27
        if isinstance(text, bytes):
28
            return text
29
        return [ord(c)for c in text]
30

31
    def _b2s(binary):
32
        return bytes(binary)
33
else:
34
    def bytes(s, e): return s
35

36

37
def _gmul(a, b):
38
    r = 0
39
    while b:
40
        if b & 1:
41
            r ^= a
42
        a <<= 1
43
        if a > 255:
44
            a ^= 0x11B
45
        b >>= 1
46
    return r
47

48

49
def _mix(n, vec):
50
    return sum(_gmul(n, v) << (24 - 8 * shift) for shift, v in enumerate(vec))
51

52

53
def _ror32(n):
54
    return (n & 255) << 24 | n >> 8
55

56

57
def _rcon():
58
    return [_gmul(1, 1 << n) for n in range(30)]
59

60

61
def _Si(S):
62
    return [S.index(n) for n in range(len(S))]
63

64

65
def _mixl(S, vec):
66
    return [_mix(s, vec) for s in S]
67

68

69
def _rorl(T):
70
    return [_ror32(t) for t in T]
71

72

73
empty = struct.pack('')
74

75

76
class AESCBC(object):
77
    nrs = {16: 10, 24: 12, 32: 14}
78
    rcon = _rcon()
79
    S = [
80
        99, 124, 119, 123, 242, 107, 111, 197, 48, 1, 103, 43, 254, 215, 171,
81
        118, 202, 130, 201, 125, 250, 89, 71, 240, 173, 212, 162, 175, 156,
82
        164, 114, 192, 183, 253, 147, 38, 54, 63, 247, 204, 52, 165, 229, 241,
83
        113, 216, 49, 21, 4, 199, 35, 195, 24, 150, 5, 154, 7, 18, 128, 226,
84
        235, 39, 178, 117, 9, 131, 44, 26, 27, 110, 90, 160, 82, 59, 214, 179,
85
        41, 227, 47, 132, 83, 209, 0, 237, 32, 252, 177, 91, 106, 203, 190, 57,
86
        74, 76, 88, 207, 208, 239, 170, 251, 67, 77, 51, 133, 69, 249, 2, 127,
87
        80, 60, 159, 168, 81, 163, 64, 143, 146, 157, 56, 245, 188, 182, 218,
88
        33, 16, 255, 243, 210, 205, 12, 19, 236, 95, 151, 68, 23, 196, 167,
89
        126, 61, 100, 93, 25, 115, 96, 129, 79, 220, 34, 42, 144, 136, 70, 238,
90
        184, 20, 222, 94, 11, 219, 224, 50, 58, 10, 73, 6, 36, 92, 194, 211,
91
        172, 98, 145, 149, 228, 121, 231, 200, 55, 109, 141, 213, 78, 169, 108,
92
        86, 244, 234, 101, 122, 174, 8, 186, 120, 37, 46, 28, 166, 180, 198,
93
        232, 221, 116, 31, 75, 189, 139, 138, 112, 62, 181, 102, 72, 3, 246,
94
        14, 97, 53, 87, 185, 134, 193, 29, 158, 225, 248, 152, 17, 105, 217,
95
        142, 148, 155, 30, 135, 233, 206, 85, 40, 223, 140, 161, 137, 13, 191,
96
        230, 66, 104, 65, 153, 45, 15, 176, 84, 187, 22
97
    ]
98
    Si = _Si(S)
99
    T1 = _mixl(S, (2, 1, 1, 3))
100
    T2 = _rorl(T1)
101
    T3 = _rorl(T2)
102
    T4 = _rorl(T3)
103
    T5 = _mixl(Si, (14, 9, 13, 11))
104
    T6 = _rorl(T5)
105
    T7 = _rorl(T6)
106
    T8 = _rorl(T7)
107
    U1 = _mixl(range(256), (14, 9, 13, 11))
108
    U2 = _rorl(U1)
109
    U3 = _rorl(U2)
110
    U4 = _rorl(U3)
111

112
    def __init__(self, key):
113
        if len(key)not in (16, 24, 32):
114
            raise ValueError('Invalid key size')
115
        rds = self.nrs[len(key)]
116
        self._Ke = [[0] * 4 for i in xrange(rds + 1)]
117
        self._Kd = [[0] * 4 for i in xrange(rds + 1)]
118
        rnd_kc = (rds + 1) * 4
119
        KC = len(key) // 4
120
        tk = [struct.unpack('>i', key[i:i + 4])[0]
121
              for i in xrange(0, len(key), 4)]
122
        rconpointer = 0
123
        t = KC
124
        for i in xrange(0, KC):
125
            self._Ke[i // 4][i % 4] = tk[i]
126
            self._Kd[rds - (i // 4)][i % 4] = tk[i]
127
        while t < rnd_kc:
128
            tt = tk[KC - 1]
129
            tk[0] ^= ((self.S[(tt >> 16) & 255] << 24) ^ (self.S[(tt >> 8) & 255] << 16) ^ (
130
                self.S[tt & 255] << 8) ^ self.S[(tt >> 24) & 255] ^ (self.rcon[rconpointer] << 24))
131
            rconpointer += 1
132
            if KC != 8:
133
                for i in xrange(1, KC):
134
                    tk[i] ^= tk[i - 1]
135
            else:
136
                for i in xrange(1, KC // 2):
137
                    tk[i] ^= tk[i - 1]
138
                tt = tk[KC // 2 - 1]
139
                tk[KC // 2] ^= (self.S[tt & 255] ^ (self.S[(tt >> 8) & 255] << 8) ^
140
                                (self.S[(tt >> 16) & 255] << 16) ^ (self.S[(tt >> 24) & 255] << 24))
141
                for i in xrange(KC // 2 + 1, KC):
142
                    tk[i] ^= tk[i - 1]
143
            j = 0
144
            while j < KC and t < rnd_kc:
145
                self._Ke[t // 4][t % 4] = tk[j]
146
                self._Kd[rds - (t // 4)][t % 4] = tk[j]
147
                j += 1
148
                t += 1
149
        for r in xrange(1, rds):
150
            for j in xrange(0, 4):
151
                tt = self._Kd[r][j]
152
                self._Kd[r][j] = (self.U1[(tt >> 24) & 255] ^ self.U2[(
153
                    tt >> 16) & 255] ^ self.U3[(tt >> 8) & 255] ^ self.U4[tt & 255])
154

155
    def _encdec(self, data, K, s, S, L1, L2, L3, L4):
156
        if len(data) != 16:
157
            raise ValueError('wrong block length')
158
        rds = len(K) - 1
159
        (s1, s2, s3) = s
160
        a = [0, 0, 0, 0]
161
        t = [(_cw(data[4 * i:4 * i + 4]) ^ K[0][i])for i in xrange(0, 4)]
162
        for r in xrange(1, rds):
163
            for i in xrange(0, 4):
164
                a[i] = L1[(t[i] >> 24) & 255]
165
                a[i] ^= L2[(t[(i + s1) % 4] >> 16) & 255]
166
                a[i] ^= L3[(t[(i + s2) % 4] >> 8) & 255]
167
                a[i] ^= L4[t[(i + s3) % 4] & 255] ^ K[r][i]
168
            t = copy.copy(a)
169
        rst = []
170
        for i in xrange(0, 4):
171
            tt = K[rds][i]
172
            rst.append((S[(t[i] >> 24) & 255] ^ (tt >> 24)) & 255)
173
            rst.append((S[(t[(i + s1) % 4] >> 16) & 255] ^ (tt >> 16)) & 255)
174
            rst.append((S[(t[(i + s2) % 4] >> 8) & 255] ^ (tt >> 8)) & 255)
175
            rst.append((S[t[(i + s3) % 4] & 255] ^ tt) & 255)
176
        return rst
177

178
    def enc_in(self, pt):
179
        return self._encdec(
180
            pt, self._Ke, [
181
                1, 2, 3], self.S, self.T1, self.T2, self.T3, self.T4)
182

183
    def dec_in(self, ct):
184
        return self._encdec(
185
            ct, self._Kd, [
186
                3, 2, 1], self.Si, self.T5, self.T6, self.T7, self.T8)
187

188
    def pad(self, pt):
189
        c = 16 - (len(pt) % 16)
190
        return pt + bytes(chr(c) * c, 'utf-8')
191

192
    def unpad(self, pt):
193
        c = pt[-1]
194
        if not isinstance(c, int):
195
            c = ord(c)
196
        return pt[:-c]
197

198
    def encrypt(self, iv, pt):
199
        if len(iv) != 16:
200
            raise ValueError('initialization vector must be 16 bytes')
201
        else:
202
            self._lcb = _s2b(iv)
203
        pt = self.pad(pt)
204
        return empty.join([self.enc_b(b)for b in chunks(pt, 16)])
205

206
    def enc_b(self, pt):
207
        if len(pt) != 16:
208
            raise ValueError('plaintext block must be 16 bytes')
209
        pt = _s2b(pt)
210
        pcb = [(p ^ l)for (p, l) in zip(pt, self._lcb)]
211
        self._lcb = self.enc_in(pcb)
212
        return _b2s(self._lcb)
213

214
    def decrypt(self, iv, ct):
215
        if len(iv) != 16:
216
            raise ValueError('initialization vector must be 16 bytes')
217
        else:
218
            self._lcb = _s2b(iv)
219
        if len(ct) % 16 != 0:
220
            raise ValueError('ciphertext must be a multiple of 16')
221
        return self.unpad(empty.join([self.dec_b(b)for b in chunks(ct, 16)]))
222

223
    def dec_b(self, ct):
224
        if len(ct) != 16:
225
            raise ValueError('ciphertext block must be 16 bytes')
226
        cb = _s2b(ct)
227
        pt = [(p ^ l)for (p, l) in zip(self.dec_in(cb), self._lcb)]
228
        self._lcb = cb
229
        return _b2s(pt)
230

231