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         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 pegs 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 peg_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_peg_letters, ring_setting=1, position='a', raw_transform=False):
 229         self.ring_peg_letters = ring_peg_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.peg_positions = [(pos(p) - pos(position)) % 26  for p in self.ring_peg_letters]
 246         self.peg_positions = [(pos(p) - (self.position + self.ring_setting - 1)) % 26  for p in self.ring_peg_letters]
 247
 248     def advance(self):
 249         super(Wheel, self).advance()
 250         self.peg_positions = [(p - 1) % 26 for p in self.peg_positions]
 251
 252
 253 class Enigma(object):
 254     """An Enigma machine.
 255
 256
 257     """
 258     def __init__(self, reflector_spec,
 259                  left_wheel_spec, left_wheel_pegs,
 260                  middle_wheel_spec, middle_wheel_pegs,
 261                  right_wheel_spec, right_wheel_pegs,
 262                  left_ring_setting, middle_ring_setting, right_ring_setting,
 263                  plugboard_setting):
 264         self.reflector = Reflector(reflector_spec)
 265         self.left_wheel = Wheel(left_wheel_spec, left_wheel_pegs, ring_setting=left_ring_setting)
 266         self.middle_wheel = Wheel(middle_wheel_spec, middle_wheel_pegs, ring_setting=middle_ring_setting)
 267         self.right_wheel = Wheel(right_wheel_spec, right_wheel_pegs, ring_setting=right_ring_setting)
 268         self.plugboard = Plugboard(plugboard_setting)
 269
 270     def __getattribute__(self,name):
 271         if name=='wheel_positions':
 272             return self.left_wheel.position, self.middle_wheel.position, self.right_wheel.position
 273         elif name=='wheel_positions_l':
 274             return self.left_wheel.position_l, self.middle_wheel.position_l, self.right_wheel.position_l
 275         elif name=='peg_positions':
 276             return self.left_wheel.peg_positions, self.middle_wheel.peg_positions, self.right_wheel.peg_positions
 277         else:
 278             return object.__getattribute__(self, name)
 279
 280     def set_wheels(self, left_wheel_position, middle_wheel_position, right_wheel_position):
 281         self.left_wheel.set_position(left_wheel_position)
 282         self.middle_wheel.set_position(middle_wheel_position)
 283         self.right_wheel.set_position(right_wheel_position)
 284
 285     def lookup(self, letter):
 286         a = self.plugboard.forward(letter)
 287         b = self.right_wheel.forward(a)
 288         c = self.middle_wheel.forward(b)
 289         d = self.left_wheel.forward(c)
 290         e = self.reflector.forward(d)
 291         f = self.left_wheel.backward(e)
 292         g = self.middle_wheel.backward(f)
 293         h = self.right_wheel.backward(g)
 294         i = self.plugboard.backward(h)
 295         return i
 296
 297     def advance(self):
 298         advance_middle = False
 299         advance_left = False
 300         if 0 in self.right_wheel.peg_positions:
 301             advance_middle = True
 302         if 0 in self.middle_wheel.peg_positions:
 303             advance_left = True
 304             advance_middle = True
 305         self.right_wheel.advance()
 306         if advance_middle: self.middle_wheel.advance()
 307         if advance_left: self.left_wheel.advance()
 308
 309     def encipher_letter(self, letter):
 310         self.advance()
 311         return self.lookup(letter)
 312
 313     def encipher(self, message):
 314         enciphered = ''
 315         for letter in clean(message):
 316             enciphered += self.encipher_letter(letter)
 317         return enciphered
 318
 319     decipher = encipher
 320
 321
 322 # for i in range(26):
 323 #     enigma.advance()
 324 #     print('enigma.advance()')
 325 #     print("assert(enigma.wheel_positions == {})".format(enigma.wheel_positions))
 326 #     print("assert(cat(enigma.wheel_positions_l) == '{}')".format(cat(enigma.wheel_positions_l)))
 327 #     print("assert(enigma.peg_positions == {})".format(enigma.peg_positions))
 328 #     print("assert(cat(enigma.lookup(l) for l in string.ascii_lowercase) == '{}')".format(cat(enigma.lookup(l) for l in string.ascii_lowercase)))
 329 #     print()
 330
 331
 332 ##################################
 333 # # Bombe
 334 ##################################
 335 #
 336 # Good explanation of [how the bombe worked](http://www.ellsbury.com/enigmabombe.htm) by Graham Ellsbury
 337 #
 338
 339 Signal = collections.namedtuple('Signal', ['bank', 'wire'])
 340 Connection = collections.namedtuple('Connection', ['banks', 'scrambler'])
 341 MenuItem = collections.namedtuple('MenuIem', ['before', 'after', 'number'])
 342
 343
 344 class Scrambler(object):
 345     def __init__(self, wheel1_spec, wheel2_spec, wheel3_spec, reflector_spec,
 346                  wheel1_pos='a', wheel2_pos='a', wheel3_pos='a'):
 347         self.wheel1 = SimpleWheel(wheel1_spec, position=wheel1_pos)
 348         self.wheel2 = SimpleWheel(wheel2_spec, position=wheel2_pos)
 349         self.wheel3 = SimpleWheel(wheel3_spec, position=wheel3_pos)
 350         self.reflector = Reflector(reflector_spec)
 351
 352     def __getattribute__(self, name):
 353         if name=='wheel_positions':
 354             return self.wheel1.position, self.wheel2.position, self.wheel3.position
 355         elif name=='wheel_positions_l':
 356             return self.wheel1.position_l, self.wheel2.position_l, self.wheel3.position_l
 357         else:
 358             return object.__getattribute__(self, name)
 359
 360     def advance(self, wheel1=False, wheel2=False, wheel3=True):
 361         if wheel1: self.wheel1.advance()
 362         if wheel2: self.wheel2.advance()
 363         if wheel3: self.wheel3.advance()
 364
 365     def lookup(self, letter):
 366         a = self.wheel3.forward(letter)
 367         b = self.wheel2.forward(a)
 368         c = self.wheel1.forward(b)
 369         d = self.reflector.forward(c)
 370         e = self.wheel1.backward(d)
 371         f = self.wheel2.backward(e)
 372         g = self.wheel3.backward(f)
 373         return g
 374
 375     def set_positions(self, wheel1_pos, wheel2_pos, wheel3_pos):
 376         self.wheel1.set_position(wheel1_pos)
 377         self.wheel2.set_position(wheel2_pos)
 378         self.wheel3.set_position(wheel3_pos)
 379
 380
 381 class Bombe(object):
 382     def __init__(self, wheel1_spec, wheel2_spec, wheel3_spec, reflector_spec,
 383                 menu=None, start_signal=None, use_diagonal_board=True,
 384                 verify_plugboard=True):
 385         self.connections = []
 386         self.wheel1_spec = wheel1_spec
 387         self.wheel2_spec = wheel2_spec
 388         self.wheel3_spec = wheel3_spec
 389         self.reflector_spec = reflector_spec
 390         if menu:
 391             self.read_menu(menu)
 392         if start_signal:
 393             self.test_start = start_signal
 394         self.use_diagonal_board = use_diagonal_board
 395         self.verify_plugboard = verify_plugboard
 396
 397     def __getattribute__(self, name):
 398         if name=='wheel_positions':
 399             return self.connections[0].scrambler.wheel_positions
 400         elif name=='wheel_positions_l':
 401             return self.connections[0].scrambler.wheel_positions_l
 402         else:
 403             return object.__getattribute__(self, name)
 404
 405     def __call__(self, start_positions):
 406         return start_positions, self.test(initial_signal=self.test_start,
 407             start_positions=start_positions,
 408             use_diagonal_board=self.use_diagonal_board,
 409             verify_plugboard=self.verify_plugboard)
 410
 411     def add_connection(self, bank_before, bank_after, scrambler):
 412         self.connections += [Connection([bank_before, bank_after], scrambler)]
 413
 414     def read_menu(self, menu):
 415         for item in menu:
 416             scrambler = Scrambler(self.wheel1_spec, self.wheel2_spec, self.wheel3_spec,
 417                                   self.reflector_spec,
 418                                   wheel3_pos=unpos(item.number - 1))
 419             self.add_connection(item.before, item.after, scrambler)
 420         most_common_letter = (collections.Counter(m.before for m in menu) + \
 421                               collections.Counter(m.after for m in menu)).most_common(1)[0][0]
 422         self.test_start = Signal(most_common_letter, most_common_letter)
 423
 424     def set_positions(self, wheel1_pos, wheel2_pos, wheel3_pos):
 425         for i, c in enumerate(self.connections):
 426             c.scrambler.set_positions(wheel1_pos, wheel2_pos, unpos(pos(wheel3_pos) + i))
 427
 428     def test(self, initial_signal=None, start_positions=None, use_diagonal_board=True,
 429             verify_plugboard=True):
 430         self.banks = {label:
 431                       dict(zip(string.ascii_lowercase, [False]*len(string.ascii_lowercase)))
 432                       for label in string.ascii_lowercase}
 433         if start_positions:
 434             self.set_positions(*start_positions)
 435         if not initial_signal:
 436             initial_signal = self.test_start
 437         self.pending = [initial_signal]
 438         self.propagate(use_diagonal_board)
 439         live_wire_count = len([self.banks[self.test_start.bank][w]
 440                     for w in self.banks[self.test_start.bank]
 441                     if self.banks[self.test_start.bank][w]])
 442         if live_wire_count < 26:
 443             if verify_plugboard:
 444                 possibles = self.possible_plugboards()
 445                 return all(s0.isdisjoint(s1) for s0 in possibles for s1 in possibles if s0 != s1)
 446             else:
 447                 return True
 448         else:
 449             return False
 450
 451     def propagate(self, use_diagonal_board):
 452         while self.pending:
 453             current = self.pending[0]
 454             # print("processing", current)
 455             self.pending = self.pending[1:]
 456             if not self.banks[current.bank][current.wire]:
 457                 self.banks[current.bank][current.wire] = True
 458                 if use_diagonal_board:
 459                     self.pending += [Signal(current.wire, current.bank)]
 460                 for c in self.connections:
 461                     if current.bank in c.banks:
 462                         other_bank = [b for b in c.banks if b != current.bank][0]
 463                         other_wire = c.scrambler.lookup(current.wire)
 464                         # print("  adding", other_bank, other_wire, "because", c.banks)
 465                         self.pending += [Signal(other_bank, other_wire)]
 466
 467     def run(self, run_start=None, wheel1_pos='a', wheel2_pos='a', wheel3_pos='a', use_diagonal_board=True):
 468         if not run_start:
 469             run_start = self.test_start
 470         self.solutions = []
 471         self.set_positions(wheel1_pos, wheel2_pos, wheel3_pos)
 472         for run_index in range(26*26*26):
 473             if self.test(initial_signal=run_start, use_diagonal_board=use_diagonal_board):
 474                 self.solutions += [self.connections[0].scrambler.wheel_positions_l]
 475             advance3 = True
 476             advance2 = False
 477             advance1 = False
 478             if (run_index + 1) % 26 == 0: advance2 = True
 479             if (run_index + 1) % (26*26) == 0: advance1 = True
 480             for c in self.connections:
 481                 c.scrambler.advance(advance1, advance2, advance3)
 482         return self.solutions
 483
 484     def possible_plugboards(self):
 485         possibles = set()
 486         for b in self.banks:
 487             active = [w for w in self.banks[b] if self.banks[b][w]]
 488             inactive = [w for w in self.banks[b] if not self.banks[b][w]]
 489             if len(active) == 1:
 490                 possibles = possibles.union({frozenset((b, active[0]))})
 491             if len(inactive) == 1:
 492                 possibles = possibles.union({frozenset((b, inactive[0]))})
 493         return possibles
 494
 495
 496 def make_menu(plaintext, ciphertext):
 497     return [MenuItem(p, c, i+1)
 498             for i, (p, c) in enumerate(zip(plaintext, ciphertext))]
 499
 500
 501 def run_multi_bombe(wheel1_spec, wheel2_spec, wheel3_spec, reflector_spec, menu,
 502                     start_signal=None, use_diagonal_board=True,
 503                     verify_plugboard=True):
 504     allwheels = itertools.product(string.ascii_lowercase, repeat=3)
 505
 506     with multiprocessing.Pool() as pool:
 507         res = pool.map(Bombe(wheel1_spec, wheel2_spec, wheel3_spec,
 508             reflector_spec, menu=menu, start_signal=start_signal,
 509             use_diagonal_board=use_diagonal_board,
 510             verify_plugboard=verify_plugboard),
 511                   allwheels)
 512     return [r[0] for r in res if r[1]]
 513
 514
 515 if __name__ == "__main__":
 516     import doctest
 517     # doctest.testmod(extraglobs={'lt': LetterTransformer(1, 'a')})
 518     doctest.testmod()
 519