daec253af300f6e39e23288b9609f781f1c97322
[szyfrow.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_notches = ['q']
49 wheel_ii_notches = ['e']
50 wheel_iii_notches = ['v']
51 wheel_iv_notches = ['j']
52 wheel_v_notches = ['z']
53 wheel_vi_notches = ['z', 'm']
54 wheel_vii_notches = ['z', 'm']
55 wheel_viii_notches = ['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 notches 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 notch_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_notch_letters, ring_setting=1, position='a', raw_transform=False):
229 self.ring_notch_letters = ring_notch_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.notch_positions = [(pos(p) - pos(position)) % 26 for p in self.ring_notch_letters]
246 # self.notch_positions = [(pos(p) - (self.position + self.ring_setting - 1)) % 26 for p in self.ring_notch_letters]
247 self.notch_positions = [(self.position + self.ring_setting - 1 - pos(p)) % 26 for p in self.ring_notch_letters]
248
249 def advance(self):
250 super(Wheel, self).advance()
251 self.notch_positions = [(p + 1) % 26 for p in self.notch_positions]
252 return self.position
253
254
255 class Enigma(object):
256 """An Enigma machine.
257
258
259 """
260 def __init__(self, reflector_spec,
261 left_wheel_spec, left_wheel_notches,
262 middle_wheel_spec, middle_wheel_notches,
263 right_wheel_spec, right_wheel_notches,
264 left_ring_setting, middle_ring_setting, right_ring_setting,
265 plugboard_setting):
266 self.reflector = Reflector(reflector_spec)
267 self.left_wheel = Wheel(left_wheel_spec, left_wheel_notches, ring_setting=left_ring_setting)
268 self.middle_wheel = Wheel(middle_wheel_spec, middle_wheel_notches, ring_setting=middle_ring_setting)
269 self.right_wheel = Wheel(right_wheel_spec, right_wheel_notches, ring_setting=right_ring_setting)
270 self.plugboard = Plugboard(plugboard_setting)
271
272 def __getattribute__(self,name):
273 if name=='wheel_positions':
274 return self.left_wheel.position, self.middle_wheel.position, self.right_wheel.position
275 elif name=='wheel_positions_l':
276 return self.left_wheel.position_l, self.middle_wheel.position_l, self.right_wheel.position_l
277 elif name=='notch_positions':
278 return self.left_wheel.notch_positions, self.middle_wheel.notch_positions, self.right_wheel.notch_positions
279 else:
280 return object.__getattribute__(self, name)
281
282 def set_wheels(self, left_wheel_position, middle_wheel_position, right_wheel_position):
283 self.left_wheel.set_position(left_wheel_position)
284 self.middle_wheel.set_position(middle_wheel_position)
285 self.right_wheel.set_position(right_wheel_position)
286
287 def lookup(self, letter):
288 a = self.plugboard.forward(letter)
289 b = self.right_wheel.forward(a)
290 c = self.middle_wheel.forward(b)
291 d = self.left_wheel.forward(c)
292 e = self.reflector.forward(d)
293 f = self.left_wheel.backward(e)
294 g = self.middle_wheel.backward(f)
295 h = self.right_wheel.backward(g)
296 i = self.plugboard.backward(h)
297 return i
298
299 def advance(self):
300 advance_middle = False
301 advance_left = False
302 if 0 in self.right_wheel.notch_positions:
303 advance_middle = True
304 if 0 in self.middle_wheel.notch_positions:
305 advance_left = True
306 advance_middle = True
307 self.right_wheel.advance()
308 if advance_middle: self.middle_wheel.advance()
309 if advance_left: self.left_wheel.advance()
310
311 def encipher_letter(self, letter):
312 self.advance()
313 return self.lookup(letter)
314
315 def encipher(self, message):
316 enciphered = ''
317 for letter in clean(message):
318 enciphered += self.encipher_letter(letter)
319 return enciphered
320
321 decipher = encipher
322
323
324 # for i in range(26):
325 # enigma.advance()
326 # print('enigma.advance()')
327 # print("assert(enigma.wheel_positions == {})".format(enigma.wheel_positions))
328 # print("assert(cat(enigma.wheel_positions_l) == '{}')".format(cat(enigma.wheel_positions_l)))
329 # print("assert(enigma.notch_positions == {})".format(enigma.notch_positions))
330 # print("assert(cat(enigma.lookup(l) for l in string.ascii_lowercase) == '{}')".format(cat(enigma.lookup(l) for l in string.ascii_lowercase)))
331 # print()
332
333
334 if __name__ == "__main__":
335 import doctest
336 # doctest.testmod(extraglobs={'lt': LetterTransformer(1, 'a')})
337 doctest.testmod()
338