tlslite.utils.rijndael

Package tlslite :: Package utils :: Module rijndael |
[hide private]
[frames] | no frames]

# Source Code for Module tlslite.utils.rijndael

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