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