qwebirc/util/rijndael.py

   1 """
   2 A pure python (slow) implementation of rijndael with a decent interface
   3
   4 To include -
   5
   6 from rijndael import rijndael
   7
   8 To do a key setup -
   9
  10 r = rijndael(key, block_size = 16)
  11
  12 key must be a string of length 16, 24, or 32
  13 blocksize must be 16, 24, or 32. Default is 16
  14
  15 To use -
  16
  17 ciphertext = r.encrypt(plaintext)
  18 plaintext = r.decrypt(ciphertext)
  19
  20 If any strings are of the wrong length a ValueError is thrown
  21 """
  22
  23 # ported from the Java reference code by Bram Cohen, April 2001
  24 # this code is public domain, unless someone makes
  25 # an intellectual property claim against the reference
  26 # code, in which case it can be made public domain by
  27 # deleting all the comments and renaming all the variables
  28
  29 import copy
  30 import string
  31
  32 shifts = [[[0, 0], [1, 3], [2, 2], [3, 1]],
  33           [[0, 0], [1, 5], [2, 4], [3, 3]],
  34           [[0, 0], [1, 7], [3, 5], [4, 4]]]
  35
  36 # [keysize][block_size]
  37 num_rounds = {16: {16: 10, 24: 12, 32: 14}, 24: {16: 12, 24: 12, 32: 14}, 32: {16: 14, 24: 14, 32: 14}}
  38
  39 A = [[1, 1, 1, 1, 1, 0, 0, 0],
  40      [0, 1, 1, 1, 1, 1, 0, 0],
  41      [0, 0, 1, 1, 1, 1, 1, 0],
  42      [0, 0, 0, 1, 1, 1, 1, 1],
  43      [1, 0, 0, 0, 1, 1, 1, 1],
  44      [1, 1, 0, 0, 0, 1, 1, 1],
  45      [1, 1, 1, 0, 0, 0, 1, 1],
  46      [1, 1, 1, 1, 0, 0, 0, 1]]
  47
  48 # produce log and alog tables, needed for multiplying in the
  49 # field GF(2^m) (generator = 3)
  50 alog = [1]
  51 for i in xrange(255):
  52     j = (alog[-1] << 1) ^ alog[-1]
  53     if j & 0x100 != 0:
  54         j ^= 0x11B
  55     alog.append(j)
  56
  57 log = [0] * 256
  58 for i in xrange(1, 255):
  59     log[alog[i]] = i
  60
  61 # multiply two elements of GF(2^m)
  62 def mul(a, b):
  63     if a == 0 or b == 0:
  64         return 0
  65     return alog[(log[a & 0xFF] + log[b & 0xFF]) % 255]
  66
  67 # substitution box based on F^{-1}(x)
  68 box = [[0] * 8 for i in xrange(256)]
  69 box[1][7] = 1
  70 for i in xrange(2, 256):
  71     j = alog[255 - log[i]]
  72     for t in xrange(8):
  73         box[i][t] = (j >> (7 - t)) & 0x01
  74
  75 B = [0, 1, 1, 0, 0, 0, 1, 1]
  76
  77 # affine transform:  box[i] <- B + A*box[i]
  78 cox = [[0] * 8 for i in xrange(256)]
  79 for i in xrange(256):
  80     for t in xrange(8):
  81         cox[i][t] = B[t]
  82         for j in xrange(8):
  83             cox[i][t] ^= A[t][j] * box[i][j]
  84
  85 # S-boxes and inverse S-boxes
  86 S =  [0] * 256
  87 Si = [0] * 256
  88 for i in xrange(256):
  89     S[i] = cox[i][0] << 7
  90     for t in xrange(1, 8):
  91         S[i] ^= cox[i][t] << (7-t)
  92     Si[S[i] & 0xFF] = i
  93
  94 # T-boxes
  95 G = [[2, 1, 1, 3],
  96     [3, 2, 1, 1],
  97     [1, 3, 2, 1],
  98     [1, 1, 3, 2]]
  99
 100 AA = [[0] * 8 for i in xrange(4)]
 101
 102 for i in xrange(4):
 103     for j in xrange(4):
 104         AA[i][j] = G[i][j]
 105         AA[i][i+4] = 1
 106
 107 for i in xrange(4):
 108     pivot = AA[i][i]
 109     if pivot == 0:
 110         t = i + 1
 111         while AA[t][i] == 0 and t < 4:
 112             t += 1
 113             assert t != 4, 'G matrix must be invertible'
 114             for j in xrange(8):
 115                 AA[i][j], AA[t][j] = AA[t][j], AA[i][j]
 116             pivot = AA[i][i]
 117     for j in xrange(8):
 118         if AA[i][j] != 0:
 119             AA[i][j] = alog[(255 + log[AA[i][j] & 0xFF] - log[pivot & 0xFF]) % 255]
 120     for t in xrange(4):
 121         if i != t:
 122             for j in xrange(i+1, 8):
 123                 AA[t][j] ^= mul(AA[i][j], AA[t][i])
 124             AA[t][i] = 0
 125
 126 iG = [[0] * 4 for i in xrange(4)]
 127
 128 for i in xrange(4):
 129     for j in xrange(4):
 130         iG[i][j] = AA[i][j + 4]
 131
 132 def mul4(a, bs):
 133     if a == 0:
 134         return 0
 135     r = 0
 136     for b in bs:
 137         r <<= 8
 138         if b != 0:
 139             r = r | mul(a, b)
 140     return r
 141
 142 T1 = []
 143 T2 = []
 144 T3 = []
 145 T4 = []
 146 T5 = []
 147 T6 = []
 148 T7 = []
 149 T8 = []
 150 U1 = []
 151 U2 = []
 152 U3 = []
 153 U4 = []
 154
 155 for t in xrange(256):
 156     s = S[t]
 157     T1.append(mul4(s, G[0]))
 158     T2.append(mul4(s, G[1]))
 159     T3.append(mul4(s, G[2]))
 160     T4.append(mul4(s, G[3]))
 161
 162     s = Si[t]
 163     T5.append(mul4(s, iG[0]))
 164     T6.append(mul4(s, iG[1]))
 165     T7.append(mul4(s, iG[2]))
 166     T8.append(mul4(s, iG[3]))
 167
 168     U1.append(mul4(t, iG[0]))
 169     U2.append(mul4(t, iG[1]))
 170     U3.append(mul4(t, iG[2]))
 171     U4.append(mul4(t, iG[3]))
 172
 173 # round constants
 174 rcon = [1]
 175 r = 1
 176 for t in xrange(1, 30):
 177     r = mul(2, r)
 178     rcon.append(r)
 179
 180 del A
 181 del AA
 182 del pivot
 183 del B
 184 del G
 185 del box
 186 del log
 187 del alog
 188 del i
 189 del j
 190 del r
 191 del s
 192 del t
 193 del mul
 194 del mul4
 195 del cox
 196 del iG
 197
 198 class rijndael:
 199     def __init__(self, key, block_size = 16):
 200         if block_size != 16 and block_size != 24 and block_size != 32:
 201             raise ValueError('Invalid block size: ' + str(block_size))
 202         if len(key) != 16 and len(key) != 24 and len(key) != 32:
 203             raise ValueError('Invalid key size: ' + str(len(key)))
 204         self.block_size = block_size
 205
 206         ROUNDS = num_rounds[len(key)][block_size]
 207         BC = block_size / 4
 208         # encryption round keys
 209         Ke = [[0] * BC for i in xrange(ROUNDS + 1)]
 210         # decryption round keys
 211         Kd = [[0] * BC for i in xrange(ROUNDS + 1)]
 212         ROUND_KEY_COUNT = (ROUNDS + 1) * BC
 213         KC = len(key) / 4
 214
 215         # copy user material bytes into temporary ints
 216         tk = []
 217         for i in xrange(0, KC):
 218             tk.append((ord(key[i * 4]) << 24) | (ord(key[i * 4 + 1]) << 16) |
 219                 (ord(key[i * 4 + 2]) << 8) | ord(key[i * 4 + 3]))
 220
 221         # copy values into round key arrays
 222         t = 0
 223         j = 0
 224         while j < KC and t < ROUND_KEY_COUNT:
 225             Ke[t / BC][t % BC] = tk[j]
 226             Kd[ROUNDS - (t / BC)][t % BC] = tk[j]
 227             j += 1
 228             t += 1
 229         tt = 0
 230         rconpointer = 0
 231         while t < ROUND_KEY_COUNT:
 232             # extrapolate using phi (the round key evolution function)
 233             tt = tk[KC - 1]
 234             tk[0] ^= (S[(tt >> 16) & 0xFF] & 0xFF) << 24 ^  \
 235                      (S[(tt >>  8) & 0xFF] & 0xFF) << 16 ^  \
 236                      (S[ tt        & 0xFF] & 0xFF) <<  8 ^  \
 237                      (S[(tt >> 24) & 0xFF] & 0xFF)       ^  \
 238                      (rcon[rconpointer]    & 0xFF) << 24
 239             rconpointer += 1
 240             if KC != 8:
 241                 for i in xrange(1, KC):
 242                     tk[i] ^= tk[i-1]
 243             else:
 244                 for i in xrange(1, KC / 2):
 245                     tk[i] ^= tk[i-1]
 246                 tt = tk[KC / 2 - 1]
 247                 tk[KC / 2] ^= (S[ tt        & 0xFF] & 0xFF)       ^ \
 248                               (S[(tt >>  8) & 0xFF] & 0xFF) <<  8 ^ \
 249                               (S[(tt >> 16) & 0xFF] & 0xFF) << 16 ^ \
 250                               (S[(tt >> 24) & 0xFF] & 0xFF) << 24
 251                 for i in xrange(KC / 2 + 1, KC):
 252                     tk[i] ^= tk[i-1]
 253             # copy values into round key arrays
 254             j = 0
 255             while j < KC and t < ROUND_KEY_COUNT:
 256                 Ke[t / BC][t % BC] = tk[j]
 257                 Kd[ROUNDS - (t / BC)][t % BC] = tk[j]
 258                 j += 1
 259                 t += 1
 260         # inverse MixColumn where needed
 261         for r in xrange(1, ROUNDS):
 262             for j in xrange(BC):
 263                 tt = Kd[r][j]
 264                 Kd[r][j] = U1[(tt >> 24) & 0xFF] ^ \
 265                            U2[(tt >> 16) & 0xFF] ^ \
 266                            U3[(tt >>  8) & 0xFF] ^ \
 267                            U4[ tt        & 0xFF]
 268         self.Ke = Ke
 269         self.Kd = Kd
 270
 271     def encrypt(self, plaintext):
 272         if len(plaintext) != self.block_size:
 273             raise ValueError('wrong block length, expected ' + str(self.block_size) + ' got ' + str(len(plaintext)))
 274         Ke = self.Ke
 275
 276         BC = self.block_size / 4
 277         ROUNDS = len(Ke) - 1
 278         if BC == 4:
 279             SC = 0
 280         elif BC == 6:
 281             SC = 1
 282         else:
 283             SC = 2
 284         s1 = shifts[SC][1][0]
 285         s2 = shifts[SC][2][0]
 286         s3 = shifts[SC][3][0]
 287         a = [0] * BC
 288         # temporary work array
 289         t = []
 290         # plaintext to ints + key
 291         for i in xrange(BC):
 292             t.append((ord(plaintext[i * 4    ]) << 24 |
 293                       ord(plaintext[i * 4 + 1]) << 16 |
 294                       ord(plaintext[i * 4 + 2]) <<  8 |
 295                       ord(plaintext[i * 4 + 3])        ) ^ Ke[0][i])
 296         # apply round transforms
 297         for r in xrange(1, ROUNDS):
 298             for i in xrange(BC):
 299                 a[i] = (T1[(t[ i           ] >> 24) & 0xFF] ^
 300                         T2[(t[(i + s1) % BC] >> 16) & 0xFF] ^
 301                         T3[(t[(i + s2) % BC] >>  8) & 0xFF] ^
 302                         T4[ t[(i + s3) % BC]        & 0xFF]  ) ^ Ke[r][i]
 303             t = copy.copy(a)
 304         # last round is special
 305         result = []
 306         for i in xrange(BC):
 307             tt = Ke[ROUNDS][i]
 308             result.append((S[(t[ i           ] >> 24) & 0xFF] ^ (tt >> 24)) & 0xFF)
 309             result.append((S[(t[(i + s1) % BC] >> 16) & 0xFF] ^ (tt >> 16)) & 0xFF)
 310             result.append((S[(t[(i + s2) % BC] >>  8) & 0xFF] ^ (tt >>  8)) & 0xFF)
 311             result.append((S[ t[(i + s3) % BC]        & 0xFF] ^  tt       ) & 0xFF)
 312         return string.join(map(chr, result), '')
 313
 314     def decrypt(self, ciphertext):
 315         if len(ciphertext) != self.block_size:
 316             raise ValueError('wrong block length, expected ' + str(self.block_size) + ' got ' + str(len(ciphertext)))
 317         Kd = self.Kd
 318
 319         BC = self.block_size / 4
 320         ROUNDS = len(Kd) - 1
 321         if BC == 4:
 322             SC = 0
 323         elif BC == 6:
 324             SC = 1
 325         else:
 326             SC = 2
 327         s1 = shifts[SC][1][1]
 328         s2 = shifts[SC][2][1]
 329         s3 = shifts[SC][3][1]
 330         a = [0] * BC
 331         # temporary work array
 332         t = [0] * BC
 333         # ciphertext to ints + key
 334         for i in xrange(BC):
 335             t[i] = (ord(ciphertext[i * 4    ]) << 24 |
 336                     ord(ciphertext[i * 4 + 1]) << 16 |
 337                     ord(ciphertext[i * 4 + 2]) <<  8 |
 338                     ord(ciphertext[i * 4 + 3])        ) ^ Kd[0][i]
 339         # apply round transforms
 340         for r in xrange(1, ROUNDS):
 341             for i in xrange(BC):
 342                 a[i] = (T5[(t[ i           ] >> 24) & 0xFF] ^
 343                         T6[(t[(i + s1) % BC] >> 16) & 0xFF] ^
 344                         T7[(t[(i + s2) % BC] >>  8) & 0xFF] ^
 345                         T8[ t[(i + s3) % BC]        & 0xFF]  ) ^ Kd[r][i]
 346             t = copy.copy(a)
 347         # last round is special
 348         result = []
 349         for i in xrange(BC):
 350             tt = Kd[ROUNDS][i]
 351             result.append((Si[(t[ i           ] >> 24) & 0xFF] ^ (tt >> 24)) & 0xFF)
 352             result.append((Si[(t[(i + s1) % BC] >> 16) & 0xFF] ^ (tt >> 16)) & 0xFF)
 353             result.append((Si[(t[(i + s2) % BC] >>  8) & 0xFF] ^ (tt >>  8)) & 0xFF)
 354             result.append((Si[ t[(i + s3) % BC]        & 0xFF] ^  tt       ) & 0xFF)
 355         return string.join(map(chr, result), '')
 356
 357 def encrypt(key, block):
 358     return rijndael(key, len(block)).encrypt(block)
 359
 360 def decrypt(key, block):
 361     return rijndael(key, len(block)).decrypt(block)
 362
 363 def test():
 364     def t(kl, bl):
 365         b = 'b' * bl
 366         r = rijndael('a' * kl, bl)
 367         assert r.decrypt(r.encrypt(b)) == b
 368     t(16, 16)
 369     t(16, 24)
 370     t(16, 32)
 371     t(24, 16)
 372     t(24, 24)
 373     t(24, 32)
 374     t(32, 16)
 375     t(32, 24)
 376     t(32, 32)