Added tests for enigma machine and bombe
[cipher-training.git] / 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