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_pegs = ['q']
  49 wheel_ii_pegs = ['e']
  50 wheel_iii_pegs = ['v']
  51 wheel_iv_pegs = ['j']
  52 wheel_v_pegs = ['z']
  53 wheel_vi_pegs = ['z', 'm']
  54 wheel_vii_pegs = ['z', 'm']
  55 wheel_viii_pegs = ['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         self.position = ord(position) - ord('a')
 186
 187     def forward(self, letter):
 188         if letter in string.ascii_lowercase:
 189             return unpos((self.forward_map[(pos(letter) + self.position) % 26] - self.position))
 190         else:
 191             return ''
 192
 193     def backward(self, letter):
 194         if letter in string.ascii_lowercase:
 195             return unpos((self.backward_map[(pos(letter) + self.position) % 26] - self.position))
 196         else:
 197             return ''
 198
 199     def advance(self):
 200         self.position = (self.position + 1) % 26
 201
 202
 203
 204 class Wheel(SimpleWheel):
 205     """A wheel with a movable ring.
 206
 207     The ring holds the letters and the pegs that turn other wheels. The core
 208     holds the wiring that does the transformation.
 209
 210     The ring position is how many steps the core is turned relative to the ring.
 211     This is one-based, so a ring setting of 1 means the core and ring are
 212     aligned.
 213
 214     The position of the wheel is the position of the core (the transforms)
 215     relative to the neutral position.
 216
 217     The position_l is the position of the ring, or what would be observed
 218     by the user of the Enigma machine.
 219
 220     The peg_positions are the number of advances of this wheel before it will
 221     advance the next wheel.
 222
 223     """
 224     def __init__(self, transform, ring_peg_letters, ring_setting=1, position='a', raw_transform=False):
 225         self.ring_peg_letters = ring_peg_letters
 226         self.ring_setting = ring_setting
 227         super(Wheel, self).__init__(transform, position=position, raw_transform=raw_transform)
 228         self.set_position(position)
 229
 230     def __getattribute__(self,name):
 231         if name=='position_l':
 232             return unpos(self.position + self.ring_setting - 1)
 233         else:
 234             return object.__getattribute__(self, name)
 235
 236     def set_position(self, position):
 237         self.position = (pos(position) - self.ring_setting + 1) % 26
 238         self.peg_positions = [(pos(p) - pos(position)) % 26  for p in self.ring_peg_letters]
 239
 240     def advance(self):
 241         super(Wheel, self).advance()
 242         self.peg_positions = [(p - 1) % 26 for p in self.peg_positions]
 243
 244
 245 class Enigma(object):
 246     """An Enigma machine.
 247
 248
 249     """
 250     def __init__(self, reflector_spec,
 251                  left_wheel_spec, left_wheel_pegs,
 252                  middle_wheel_spec, middle_wheel_pegs,
 253                  right_wheel_spec, right_wheel_pegs,
 254                  left_ring_setting, middle_ring_setting, right_ring_setting,
 255                  plugboard_setting):
 256         self.reflector = Reflector(reflector_spec)
 257         self.left_wheel = Wheel(left_wheel_spec, left_wheel_pegs, ring_setting=left_ring_setting)
 258         self.middle_wheel = Wheel(middle_wheel_spec, middle_wheel_pegs, ring_setting=middle_ring_setting)
 259         self.right_wheel = Wheel(right_wheel_spec, right_wheel_pegs, ring_setting=right_ring_setting)
 260         self.plugboard = Plugboard(plugboard_setting)
 261
 262     def __getattribute__(self,name):
 263         if name=='wheel_positions':
 264             return self.left_wheel.position, self.middle_wheel.position, self.right_wheel.position
 265         elif name=='wheel_positions_l':
 266             return self.left_wheel.position_l, self.middle_wheel.position_l, self.right_wheel.position_l
 267         elif name=='peg_positions':
 268             return self.left_wheel.peg_positions, self.middle_wheel.peg_positions, self.right_wheel.peg_positions
 269         else:
 270             return object.__getattribute__(self, name)
 271
 272     def set_wheels(self, left_wheel_position, middle_wheel_position, right_wheel_position):
 273         self.left_wheel.set_position(left_wheel_position)
 274         self.middle_wheel.set_position(middle_wheel_position)
 275         self.right_wheel.set_position(right_wheel_position)
 276
 277     def lookup(self, letter):
 278         a = self.plugboard.forward(letter)
 279         b = self.right_wheel.forward(a)
 280         c = self.middle_wheel.forward(b)
 281         d = self.left_wheel.forward(c)
 282         e = self.reflector.forward(d)
 283         f = self.left_wheel.backward(e)
 284         g = self.middle_wheel.backward(f)
 285         h = self.right_wheel.backward(g)
 286         i = self.plugboard.backward(h)
 287         return i
 288
 289     def advance(self):
 290         advance_middle = False
 291         advance_left = False
 292         if 0 in self.right_wheel.peg_positions:
 293             advance_middle = True
 294         if 0 in self.middle_wheel.peg_positions:
 295             advance_left = True
 296             advance_middle = True
 297         self.right_wheel.advance()
 298         if advance_middle: self.middle_wheel.advance()
 299         if advance_left: self.left_wheel.advance()
 300
 301     def encipher_letter(self, letter):
 302         self.advance()
 303         return self.lookup(letter)
 304
 305     def encipher(self, message):
 306         enciphered = ''
 307         for letter in clean(message):
 308             enciphered += self.encipher_letter(letter)
 309         return enciphered
 310
 311     decipher = encipher
 312
 313
 314 # for i in range(26):
 315 #     enigma.advance()
 316 #     print('enigma.advance()')
 317 #     print("assert(enigma.wheel_positions == {})".format(enigma.wheel_positions))
 318 #     print("assert(cat(enigma.wheel_positions_l) == '{}')".format(cat(enigma.wheel_positions_l)))
 319 #     print("assert(enigma.peg_positions == {})".format(enigma.peg_positions))
 320 #     print("assert(cat(enigma.lookup(l) for l in string.ascii_lowercase) == '{}')".format(cat(enigma.lookup(l) for l in string.ascii_lowercase)))
 321 #     print()
 322
 323
 324 ##################################
 325 # # Bombe
 326 ##################################
 327 #
 328 # Good explanation of [how the bombe worked](http://www.ellsbury.com/enigmabombe.htm) by Graham Ellsbury
 329 #
 330
 331 Signal = collections.namedtuple('Signal', ['bank', 'wire'])
 332 Connection = collections.namedtuple('Connection', ['banks', 'scrambler'])
 333 MenuItem = collections.namedtuple('MenuIem', ['before', 'after', 'number'])
 334
 335
 336 class Scrambler(object):
 337     def __init__(self, wheel1_spec, wheel2_spec, wheel3_spec, reflector_spec,
 338                  wheel1_pos='a', wheel2_pos='a', wheel3_pos='a'):
 339         self.wheel1 = SimpleWheel(wheel1_spec, position=wheel1_pos)
 340         self.wheel2 = SimpleWheel(wheel2_spec, position=wheel2_pos)
 341         self.wheel3 = SimpleWheel(wheel3_spec, position=wheel3_pos)
 342         self.reflector = Reflector(reflector_spec)
 343
 344     def __getattribute__(self, name):
 345         if name=='wheel_positions':
 346             return self.wheel1.position, self.wheel2.position, self.wheel3.position
 347         elif name=='wheel_positions_l':
 348             return self.wheel1.position_l, self.wheel2.position_l, self.wheel3.position_l
 349         else:
 350             return object.__getattribute__(self, name)
 351
 352     def advance(self, wheel1=False, wheel2=False, wheel3=True):
 353         if wheel1: self.wheel1.advance()
 354         if wheel2: self.wheel2.advance()
 355         if wheel3: self.wheel3.advance()
 356
 357     def lookup(self, letter):
 358         a = self.wheel3.forward(letter)
 359         b = self.wheel2.forward(a)
 360         c = self.wheel1.forward(b)
 361         d = self.reflector.forward(c)
 362         e = self.wheel1.backward(d)
 363         f = self.wheel2.backward(e)
 364         g = self.wheel3.backward(f)
 365         return g
 366
 367     def set_positions(self, wheel1_pos, wheel2_pos, wheel3_pos):
 368         self.wheel1.set_position(wheel1_pos)
 369         self.wheel2.set_position(wheel2_pos)
 370         self.wheel3.set_position(wheel3_pos)
 371
 372
 373 class Bombe(object):
 374     def __init__(self, wheel1_spec, wheel2_spec, wheel3_spec, reflector_spec,
 375                 menu=None, start_signal=None, use_diagonal_board=True,
 376                 verify_plugboard=True):
 377         self.connections = []
 378         self.wheel1_spec = wheel1_spec
 379         self.wheel2_spec = wheel2_spec
 380         self.wheel3_spec = wheel3_spec
 381         self.reflector_spec = reflector_spec
 382         if menu:
 383             self.read_menu(menu)
 384         if start_signal:
 385             self.test_start = start_signal
 386         self.use_diagonal_board = use_diagonal_board
 387         self.verify_plugboard = verify_plugboard
 388
 389     def __getattribute__(self, name):
 390         if name=='wheel_positions':
 391             return self.connections[0].scrambler.wheel_positions
 392         elif name=='wheel_positions_l':
 393             return self.connections[0].scrambler.wheel_positions_l
 394         else:
 395             return object.__getattribute__(self, name)
 396
 397     def __call__(self, start_positions):
 398         return start_positions, self.test(initial_signal=self.test_start,
 399             start_positions=start_positions,
 400             use_diagonal_board=self.use_diagonal_board,
 401             verify_plugboard=self.verify_plugboard)
 402
 403     def add_connection(self, bank_before, bank_after, scrambler):
 404         self.connections += [Connection([bank_before, bank_after], scrambler)]
 405
 406     def read_menu(self, menu):
 407         for item in menu:
 408             scrambler = Scrambler(self.wheel1_spec, self.wheel2_spec, self.wheel3_spec,
 409                                   self.reflector_spec,
 410                                   wheel3_pos=unpos(item.number - 1))
 411             self.add_connection(item.before, item.after, scrambler)
 412         most_common_letter = (collections.Counter(m.before for m in menu) + \
 413                               collections.Counter(m.after for m in menu)).most_common(1)[0][0]
 414         self.test_start = Signal(most_common_letter, most_common_letter)
 415
 416     def set_positions(self, wheel1_pos, wheel2_pos, wheel3_pos):
 417         for i, c in enumerate(self.connections):
 418             c.scrambler.set_positions(wheel1_pos, wheel2_pos, unpos(pos(wheel3_pos) + i))
 419
 420     def test(self, initial_signal=None, start_positions=None, use_diagonal_board=True,
 421             verify_plugboard=True):
 422         self.banks = {label:
 423                       dict(zip(string.ascii_lowercase, [False]*len(string.ascii_lowercase)))
 424                       for label in string.ascii_lowercase}
 425         if start_positions:
 426             self.set_positions(*start_positions)
 427         if not initial_signal:
 428             initial_signal = self.test_start
 429         self.pending = [initial_signal]
 430         self.propagate(use_diagonal_board)
 431         live_wire_count = len([self.banks[self.test_start.bank][w]
 432                     for w in self.banks[self.test_start.bank]
 433                     if self.banks[self.test_start.bank][w]])
 434         if live_wire_count < 26:
 435             if verify_plugboard:
 436                 possibles = self.possible_plugboards()
 437                 return all(s0.isdisjoint(s1) for s0 in possibles for s1 in possibles if s0 != s1)
 438             else:
 439                 return True
 440         else:
 441             return False
 442
 443     def propagate(self, use_diagonal_board):
 444         while self.pending:
 445             current = self.pending[0]
 446             # print("processing", current)
 447             self.pending = self.pending[1:]
 448             if not self.banks[current.bank][current.wire]:
 449                 self.banks[current.bank][current.wire] = True
 450                 if use_diagonal_board:
 451                     self.pending += [Signal(current.wire, current.bank)]
 452                 for c in self.connections:
 453                     if current.bank in c.banks:
 454                         other_bank = [b for b in c.banks if b != current.bank][0]
 455                         other_wire = c.scrambler.lookup(current.wire)
 456                         # print("  adding", other_bank, other_wire, "because", c.banks)
 457                         self.pending += [Signal(other_bank, other_wire)]
 458
 459     def run(self, run_start=None, wheel1_pos='a', wheel2_pos='a', wheel3_pos='a', use_diagonal_board=True):
 460         if not run_start:
 461             run_start = self.test_start
 462         self.solutions = []
 463         self.set_positions(wheel1_pos, wheel2_pos, wheel3_pos)
 464         for run_index in range(26*26*26):
 465             if self.test(initial_signal=run_start, use_diagonal_board=use_diagonal_board):
 466                 self.solutions += [self.connections[0].scrambler.wheel_positions_l]
 467             advance3 = True
 468             advance2 = False
 469             advance1 = False
 470             if (run_index + 1) % 26 == 0: advance2 = True
 471             if (run_index + 1) % (26*26) == 0: advance1 = True
 472             for c in self.connections:
 473                 c.scrambler.advance(advance1, advance2, advance3)
 474         return self.solutions
 475
 476     def possible_plugboards(self):
 477         possibles = set()
 478         for b in self.banks:
 479             active = [w for w in self.banks[b] if self.banks[b][w]]
 480             inactive = [w for w in self.banks[b] if not self.banks[b][w]]
 481             if len(active) == 1:
 482                 possibles = possibles.union({frozenset((b, active[0]))})
 483             if len(inactive) == 1:
 484                 possibles = possibles.union({frozenset((b, inactive[0]))})
 485         return possibles
 486
 487
 488 def make_menu(plaintext, ciphertext):
 489     return [MenuItem(p, c, i+1)
 490             for i, (p, c) in enumerate(zip(plaintext, ciphertext))]
 491
 492
 493 def run_multi_bombe(wheel1_spec, wheel2_spec, wheel3_spec, reflector_spec, menu,
 494                     start_signal=None, use_diagonal_board=True,
 495                     verify_plugboard=True):
 496     allwheels = itertools.product(string.ascii_lowercase, repeat=3)
 497
 498     with multiprocessing.Pool() as pool:
 499         res = pool.map(Bombe(wheel1_spec, wheel2_spec, wheel3_spec,
 500             reflector_spec, menu=menu, start_signal=start_signal,
 501             use_diagonal_board=use_diagonal_board,
 502             verify_plugboard=verify_plugboard),
 503                   allwheels)
 504     return [r[0] for r in res if r[1]]
 505
 506
 507 if __name__ == "__main__":
 508     import doctest
 509     # doctest.testmod(extraglobs={'lt': LetterTransformer(1, 'a')})
 510     doctest.testmod()
 511