enigma.py

   1
   2 # coding: utf-8
   3
   4 ##################################
   5 # # Enigma machine
   6 ##################################
   7 # Specification from [Codes and Ciphers](http://www.codesandciphers.org.uk/enigma/rotorspec.htm) page.
   8 #
   9 # Example Enigma machines from [Louise Dale](http://enigma.louisedade.co.uk/enigma.html) (full simulation) and [EnigmaCo](http://enigmaco.de/enigma/enigma.html) (good animation of the wheels, but no ring settings).
  10 #
  11 # There's also the nice Enigma simulator for Android by [Franklin Heath](https://franklinheath.co.uk/2012/02/04/our-first-app-published-enigma-simulator/), available on the [Google Play store](https://play.google.com/store/apps/details?id=uk.co.franklinheath.enigmasim&hl=en_GB).
  12
  13
  14
  15 import string
  16 import collections
  17 import multiprocessing
  18 import itertools
  19
  20 # Some convenience functions
  21
  22 cat = ''.join
  23
  24 def clean(text): return cat(l.lower() for l in text if l in string.ascii_letters)
  25
  26 def pos(letter):
  27     if letter in string.ascii_lowercase:
  28         return ord(letter) - ord('a')
  29     elif letter in string.ascii_uppercase:
  30         return ord(letter) - ord('A')
  31     else:
  32         return ''
  33
  34 def unpos(number): return chr(number % 26 + ord('a'))
  35
  36
  37 wheel_i_spec = 'ekmflgdqvzntowyhxuspaibrcj'
  38 wheel_ii_spec = 'ajdksiruxblhwtmcqgznpyfvoe'
  39 wheel_iii_spec = 'bdfhjlcprtxvznyeiwgakmusqo'
  40 wheel_iv_spec = 'esovpzjayquirhxlnftgkdcmwb'
  41 wheel_v_spec = 'vzbrgityupsdnhlxawmjqofeck'
  42 wheel_vi_spec = 'jpgvoumfyqbenhzrdkasxlictw'
  43 wheel_vii_spec = 'nzjhgrcxmyswboufaivlpekqdt'
  44 wheel_viii_spec = 'fkqhtlxocbjspdzramewniuygv'
  45 beta_wheel_spec = 'leyjvcnixwpbqmdrtakzgfuhos'
  46 gamma_wheel_spec = 'fsokanuerhmbtiycwlqpzxvgjd'
  47
  48 wheel_i_notches = ['q']
  49 wheel_ii_notches = ['e']
  50 wheel_iii_notches = ['v']
  51 wheel_iv_notches = ['j']
  52 wheel_v_notches = ['z']
  53 wheel_vi_notches = ['z', 'm']
  54 wheel_vii_notches = ['z', 'm']
  55 wheel_viii_notches = ['z', 'm']
  56
  57 reflector_b_spec = 'ay br cu dh eq fs gl ip jx kn mo tz vw'
  58 reflector_c_spec = 'af bv cp dj ei go hy kr lz mx nw tq su'
  59
  60
  61
  62 class LetterTransformer(object):
  63     """A generic substitution cipher, that has different transforms in the
  64     forward and backward directions. It requires that the transforms for all
  65     letters by provided.
  66     """
  67     def __init__(self, specification, raw_transform=False):
  68         if raw_transform:
  69             transform = specification
  70         else:
  71             transform = self.parse_specification(specification)
  72         self.validate_transform(transform)
  73         self.make_transform_map(transform)
  74
  75     def parse_specification(self, specification):
  76         return list(zip(string.ascii_lowercase, clean(specification)))
  77         # return specification
  78
  79     def validate_transform(self, transform):
  80         """A set of pairs, of from-to"""
  81         if len(transform) != 26:
  82             raise ValueError("Transform specification has {} pairs, requires 26".
  83                 format(len(transform)))
  84         for p in transform:
  85             if len(p) != 2:
  86                 raise ValueError("Not all mappings in transform "
  87                     "have two elements")
  88         if len(set([p[0] for p in transform])) != 26:
  89             raise ValueError("Transform specification must list 26 origin letters")
  90         if len(set([p[1] for p in transform])) != 26:
  91             raise ValueError("Transform specification must list 26 destination letters")
  92
  93     def make_empty_transform(self):
  94         self.forward_map = [0] * 26
  95         self.backward_map = [0] * 26
  96
  97     def make_transform_map(self, transform):
  98         self.make_empty_transform()
  99         for p in transform:
 100             self.forward_map[pos(p[0])] = pos(p[1])
 101             self.backward_map[pos(p[1])] = pos(p[0])
 102         return self.forward_map, self.backward_map
 103
 104     def forward(self, letter):
 105         if letter in string.ascii_lowercase:
 106             return unpos(self.forward_map[pos(letter)])
 107         else:
 108             return ''
 109
 110     def backward(self, letter):
 111         if letter in string.ascii_lowercase:
 112             return unpos(self.backward_map[pos(letter)])
 113         else:
 114             return ''
 115
 116
 117 class Plugboard(LetterTransformer):
 118     """A plugboard, a type of letter transformer where forward and backward
 119     transforms are the same. If a letter isn't explicitly transformed, it is
 120     kept as it is.
 121     """
 122     def parse_specification(self, specification):
 123         return [tuple(clean(p)) for p in specification.split()]
 124
 125     def validate_transform(self, transform):
 126         """A set of pairs, of from-to"""
 127         for p in transform:
 128             if len(p) != 2:
 129                 raise ValueError("Not all mappings in transform"
 130                     "have two elements")
 131
 132     def make_empty_transform(self):
 133         self.forward_map = list(range(26))
 134         self.backward_map = list(range(26))
 135
 136     def make_transform_map(self, transform):
 137         expanded_transform = transform + [tuple(reversed(p)) for p in transform]
 138         return super(Plugboard, self).make_transform_map(expanded_transform)
 139
 140
 141
 142
 143 class Reflector(Plugboard):
 144     """A reflector is a plugboard that requires 13 transforms.
 145     """
 146     def validate_transform(self, transform):
 147         if len(transform) != 13:
 148             raise ValueError("Reflector specification has {} pairs, requires 13".
 149                 format(len(transform)))
 150         if len(set([p[0] for p in transform] +
 151                     [p[1] for p in transform])) != 26:
 152             raise ValueError("Reflector specification does not contain 26 letters")
 153         try:
 154             super(Reflector, self).validate_transform(transform)
 155         except ValueError as v:
 156             raise ValueError("Not all mappings in reflector have two elements")
 157
 158
 159
 160
 161 class SimpleWheel(LetterTransformer):
 162     """A wheel is a transform that rotates.
 163
 164     Looking from the right, letters go in sequence a-b-c clockwise around the
 165     wheel.
 166
 167     The position of the wheel is the number of spaces anticlockwise the wheel
 168     has turned.
 169
 170     Letter inputs and outputs are given relative to the frame holding the wheel,
 171     so if the wheel is advanced three places, an input of 'p' will enter the
 172     wheel on the position under the wheel's 'q' label.
 173     """
 174     def __init__(self, transform, position='a', raw_transform=False):
 175         super(SimpleWheel, self).__init__(transform, raw_transform)
 176         self.set_position(position)
 177
 178     def __getattribute__(self,name):
 179         if name=='position_l':
 180             return unpos(self.position)
 181         else:
 182             return object.__getattribute__(self, name)
 183
 184     def set_position(self, position):
 185         if isinstance(position, str):
 186             # self.position = ord(position) - ord('a')
 187             self.position = pos(position)
 188         else:
 189             self.position = position
 190
 191     def forward(self, letter):
 192         if letter in string.ascii_lowercase:
 193             return unpos((self.forward_map[(pos(letter) + self.position) % 26] - self.position))
 194         else:
 195             return ''
 196
 197     def backward(self, letter):
 198         if letter in string.ascii_lowercase:
 199             return unpos((self.backward_map[(pos(letter) + self.position) % 26] - self.position))
 200         else:
 201             return ''
 202
 203     def advance(self):
 204         self.position = (self.position + 1) % 26
 205
 206
 207
 208 class Wheel(SimpleWheel):
 209     """A wheel with a movable ring.
 210
 211     The ring holds the letters and the notches that turn other wheels. The core
 212     holds the wiring that does the transformation.
 213
 214     The ring position is how many steps the core is turned relative to the ring.
 215     This is one-based, so a ring setting of 1 means the core and ring are
 216     aligned.
 217
 218     The position of the wheel is the position of the core (the transforms)
 219     relative to the neutral position.
 220
 221     The position_l is the position of the ring, or what would be observed
 222     by the user of the Enigma machine.
 223
 224     The notch_positions are the number of advances of this wheel before it will
 225     advance the next wheel.
 226
 227     """
 228     def __init__(self, transform, ring_notch_letters, ring_setting=1, position='a', raw_transform=False):
 229         self.ring_notch_letters = ring_notch_letters
 230         self.ring_setting = ring_setting
 231         super(Wheel, self).__init__(transform, position=position, raw_transform=raw_transform)
 232         self.set_position(position)
 233
 234     def __getattribute__(self,name):
 235         if name=='position_l':
 236             return unpos(self.position + self.ring_setting - 1)
 237         else:
 238             return object.__getattribute__(self, name)
 239
 240     def set_position(self, position):
 241         if isinstance(position, str):
 242             self.position = (pos(position) - self.ring_setting + 1) % 26
 243         else:
 244             self.position = (position - self.ring_setting) % 26
 245         # # self.notch_positions = [(pos(p) - pos(position)) % 26  for p in self.ring_notch_letters]
 246         # self.notch_positions = [(pos(p) - (self.position + self.ring_setting - 1)) % 26  for p in self.ring_notch_letters]
 247         self.notch_positions = [(self.position + self.ring_setting - 1 - pos(p)) % 26  for p in self.ring_notch_letters]
 248
 249     def advance(self):
 250         super(Wheel, self).advance()
 251         self.notch_positions = [(p + 1) % 26 for p in self.notch_positions]
 252         return self.position
 253
 254
 255 class Enigma(object):
 256     """An Enigma machine.
 257
 258
 259     """
 260     def __init__(self, reflector_spec,
 261                  left_wheel_spec, left_wheel_notches,
 262                  middle_wheel_spec, middle_wheel_notches,
 263                  right_wheel_spec, right_wheel_notches,
 264                  left_ring_setting, middle_ring_setting, right_ring_setting,
 265                  plugboard_setting):
 266         self.reflector = Reflector(reflector_spec)
 267         self.left_wheel = Wheel(left_wheel_spec, left_wheel_notches, ring_setting=left_ring_setting)
 268         self.middle_wheel = Wheel(middle_wheel_spec, middle_wheel_notches, ring_setting=middle_ring_setting)
 269         self.right_wheel = Wheel(right_wheel_spec, right_wheel_notches, ring_setting=right_ring_setting)
 270         self.plugboard = Plugboard(plugboard_setting)
 271
 272     def __getattribute__(self,name):
 273         if name=='wheel_positions':
 274             return self.left_wheel.position, self.middle_wheel.position, self.right_wheel.position
 275         elif name=='wheel_positions_l':
 276             return self.left_wheel.position_l, self.middle_wheel.position_l, self.right_wheel.position_l
 277         elif name=='notch_positions':
 278             return self.left_wheel.notch_positions, self.middle_wheel.notch_positions, self.right_wheel.notch_positions
 279         else:
 280             return object.__getattribute__(self, name)
 281
 282     def set_wheels(self, left_wheel_position, middle_wheel_position, right_wheel_position):
 283         self.left_wheel.set_position(left_wheel_position)
 284         self.middle_wheel.set_position(middle_wheel_position)
 285         self.right_wheel.set_position(right_wheel_position)
 286
 287     def lookup(self, letter):
 288         a = self.plugboard.forward(letter)
 289         b = self.right_wheel.forward(a)
 290         c = self.middle_wheel.forward(b)
 291         d = self.left_wheel.forward(c)
 292         e = self.reflector.forward(d)
 293         f = self.left_wheel.backward(e)
 294         g = self.middle_wheel.backward(f)
 295         h = self.right_wheel.backward(g)
 296         i = self.plugboard.backward(h)
 297         return i
 298
 299     def advance(self):
 300         advance_middle = False
 301         advance_left = False
 302         if 0 in self.right_wheel.notch_positions:
 303             advance_middle = True
 304         if 0 in self.middle_wheel.notch_positions:
 305             advance_left = True
 306             advance_middle = True
 307         self.right_wheel.advance()
 308         if advance_middle: self.middle_wheel.advance()
 309         if advance_left: self.left_wheel.advance()
 310
 311     def encipher_letter(self, letter):
 312         self.advance()
 313         return self.lookup(letter)
 314
 315     def encipher(self, message):
 316         enciphered = ''
 317         for letter in clean(message):
 318             enciphered += self.encipher_letter(letter)
 319         return enciphered
 320
 321     decipher = encipher
 322
 323
 324 # for i in range(26):
 325 #     enigma.advance()
 326 #     print('enigma.advance()')
 327 #     print("assert(enigma.wheel_positions == {})".format(enigma.wheel_positions))
 328 #     print("assert(cat(enigma.wheel_positions_l) == '{}')".format(cat(enigma.wheel_positions_l)))
 329 #     print("assert(enigma.notch_positions == {})".format(enigma.notch_positions))
 330 #     print("assert(cat(enigma.lookup(l) for l in string.ascii_lowercase) == '{}')".format(cat(enigma.lookup(l) for l in string.ascii_lowercase)))
 331 #     print()
 332
 333
 334 if __name__ == "__main__":
 335     import doctest
 336     # doctest.testmod(extraglobs={'lt': LetterTransformer(1, 'a')})
 337     doctest.testmod()
 338