2016 challenge 3
[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 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