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_pegs = ['q']
  64 wheel_ii_pegs = ['e']
  65 wheel_iii_pegs = ['v']
  66 wheel_iv_pegs = ['j']
  67 wheel_v_pegs = ['z']
  68 wheel_vi_pegs = ['z', 'm']
  69 wheel_vii_pegs = ['z', 'm']
  70 wheel_viii_pegs = ['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 pegs 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 peg_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_peg_letters, ring_setting=1, position='a', raw_transform=False):
 244         self.ring_peg_letters = ring_peg_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.peg_positions = [(pos(p) - pos(position)) % 26  for p in self.ring_peg_letters]
 261         self.peg_positions = [(pos(p) - (self.position + self.ring_setting - 1)) % 26  for p in self.ring_peg_letters]
 262
 263     def advance(self):
 264         super(Wheel, self).advance()
 265         self.peg_positions = [(p - 1) % 26 for p in self.peg_positions]
 266
 267
 268 class Enigma(object):
 269     """An Enigma machine.
 270
 271
 272     """
 273     def __init__(self, reflector_spec,
 274                  left_wheel_spec, left_wheel_pegs,
 275                  middle_wheel_spec, middle_wheel_pegs,
 276                  right_wheel_spec, right_wheel_pegs,
 277                  left_ring_setting, middle_ring_setting, right_ring_setting,
 278                  plugboard_setting):
 279         self.reflector = Reflector(reflector_spec)
 280         self.left_wheel = Wheel(left_wheel_spec, left_wheel_pegs, ring_setting=left_ring_setting)
 281         self.middle_wheel = Wheel(middle_wheel_spec, middle_wheel_pegs, ring_setting=middle_ring_setting)
 282         self.right_wheel = Wheel(right_wheel_spec, right_wheel_pegs, ring_setting=right_ring_setting)
 283         self.plugboard = Plugboard(plugboard_setting)
 284
 285     def __getattribute__(self,name):
 286         if name=='wheel_positions':
 287             return self.left_wheel.position, self.middle_wheel.position, self.right_wheel.position
 288         elif name=='wheel_positions_l':
 289             return self.left_wheel.position_l, self.middle_wheel.position_l, self.right_wheel.position_l
 290         elif name=='peg_positions':
 291             return self.left_wheel.peg_positions, self.middle_wheel.peg_positions, self.right_wheel.peg_positions
 292         else:
 293             return object.__getattribute__(self, name)
 294
 295     def set_wheels(self, left_wheel_position, middle_wheel_position, right_wheel_position):
 296         self.left_wheel.set_position(left_wheel_position)
 297         self.middle_wheel.set_position(middle_wheel_position)
 298         self.right_wheel.set_position(right_wheel_position)
 299
 300     def lookup(self, letter):
 301         a = self.plugboard.forward(letter)
 302         b = self.right_wheel.forward(a)
 303         c = self.middle_wheel.forward(b)
 304         d = self.left_wheel.forward(c)
 305         e = self.reflector.forward(d)
 306         f = self.left_wheel.backward(e)
 307         g = self.middle_wheel.backward(f)
 308         h = self.right_wheel.backward(g)
 309         i = self.plugboard.backward(h)
 310         return i
 311
 312     def advance(self):
 313         advance_middle = False
 314         advance_left = False
 315         if 0 in self.right_wheel.peg_positions:
 316             advance_middle = True
 317         if 0 in self.middle_wheel.peg_positions:
 318             advance_left = True
 319             advance_middle = True
 320         self.right_wheel.advance()
 321         if advance_middle: self.middle_wheel.advance()
 322         if advance_left: self.left_wheel.advance()
 323
 324     def encipher_letter(self, letter):
 325         self.advance()
 326         return self.lookup(letter)
 327
 328     def encipher(self, message):
 329         enciphered = ''
 330         for letter in clean(message):
 331             enciphered += self.encipher_letter(letter)
 332         return enciphered
 333
 334     decipher = encipher
 335
 336
 337 # for i in range(26):
 338 #     enigma.advance()
 339 #     print('enigma.advance()')
 340 #     print("assert(enigma.wheel_positions == {})".format(enigma.wheel_positions))
 341 #     print("assert(cat(enigma.wheel_positions_l) == '{}')".format(cat(enigma.wheel_positions_l)))
 342 #     print("assert(enigma.peg_positions == {})".format(enigma.peg_positions))
 343 #     print("assert(cat(enigma.lookup(l) for l in string.ascii_lowercase) == '{}')".format(cat(enigma.lookup(l) for l in string.ascii_lowercase)))
 344 #     print()
 345
 346
 347 if __name__ == "__main__":
 348     import doctest
 349     # doctest.testmod(extraglobs={'lt': LetterTransformer(1, 'a')})
 350     doctest.testmod()
 351