a51955acad335d9409aca193fe39b7e9eb15e487
[szyfrow.git] / szyfrow / pocket_enigma.py
1 from support.utilities import *
2 from support.language_models import *
3
4 from logger import logger
5
6
7 class PocketEnigma(object):
8 """A pocket enigma machine
9 The wheel is internally represented as a 26-element list self.wheel_map,
10 where wheel_map[i] == j shows that the position i places on from the arrow
11 maps to the position j places on.
12 """
13 def __init__(self, wheel=1, position='a'):
14 """initialise the pocket enigma, including which wheel to use and the
15 starting position of the wheel.
16
17 The wheel is either 1 or 2 (the predefined wheels) or a list of letter
18 pairs.
19
20 The position is the letter pointed to by the arrow on the wheel.
21
22 >>> pe.wheel_map
23 [25, 4, 23, 10, 1, 7, 9, 5, 12, 6, 3, 17, 8, 14, 13, 21, 19, 11, 20, 16, 18, 15, 24, 2, 22, 0]
24 >>> pe.position
25 0
26 """
27 self.wheel1 = [('a', 'z'), ('b', 'e'), ('c', 'x'), ('d', 'k'),
28 ('f', 'h'), ('g', 'j'), ('i', 'm'), ('l', 'r'), ('n', 'o'),
29 ('p', 'v'), ('q', 't'), ('s', 'u'), ('w', 'y')]
30 self.wheel2 = [('a', 'c'), ('b', 'd'), ('e', 'w'), ('f', 'i'),
31 ('g', 'p'), ('h', 'm'), ('j', 'k'), ('l', 'n'), ('o', 'q'),
32 ('r', 'z'), ('s', 'u'), ('t', 'v'), ('x', 'y')]
33 if wheel == 1:
34 self.make_wheel_map(self.wheel1)
35 elif wheel == 2:
36 self.make_wheel_map(self.wheel2)
37 else:
38 self.validate_wheel_spec(wheel)
39 self.make_wheel_map(wheel)
40 if position in string.ascii_lowercase:
41 self.position = pos(position)
42 else:
43 self.position = position
44
45 def make_wheel_map(self, wheel_spec):
46 """Expands a wheel specification from a list of letter-letter pairs
47 into a full wheel_map.
48
49 >>> pe.make_wheel_map(pe.wheel2)
50 [2, 3, 0, 1, 22, 8, 15, 12, 5, 10, 9, 13, 7, 11, 16, 6, 14, 25, 20, 21, 18, 19, 4, 24, 23, 17]
51 """
52 self.validate_wheel_spec(wheel_spec)
53 self.wheel_map = [0] * 26
54 for p in wheel_spec:
55 self.wheel_map[pos(p[0])] = pos(p[1])
56 self.wheel_map[pos(p[1])] = pos(p[0])
57 return self.wheel_map
58
59 def validate_wheel_spec(self, wheel_spec):
60 """Validates that a wheel specificaiton will turn into a valid wheel
61 map.
62
63 >>> pe.validate_wheel_spec([])
64 Traceback (most recent call last):
65 ...
66 ValueError: Wheel specification has 0 pairs, requires 13
67 >>> pe.validate_wheel_spec([('a', 'b', 'c')]*13)
68 Traceback (most recent call last):
69 ...
70 ValueError: Not all mappings in wheel specificationhave two elements
71 >>> pe.validate_wheel_spec([('a', 'b')]*13)
72 Traceback (most recent call last):
73 ...
74 ValueError: Wheel specification does not contain 26 letters
75 """
76 if len(wheel_spec) != 13:
77 raise ValueError("Wheel specification has {} pairs, requires 13".
78 format(len(wheel_spec)))
79 for p in wheel_spec:
80 if len(p) != 2:
81 raise ValueError("Not all mappings in wheel specification"
82 "have two elements")
83 if len(set([p[0] for p in wheel_spec] +
84 [p[1] for p in wheel_spec])) != 26:
85 raise ValueError("Wheel specification does not contain 26 letters")
86
87 def encipher_letter(self, letter):
88 """Enciphers a single letter, by advancing the wheel before looking up
89 the letter on the wheel.
90
91 >>> pe.set_position('f')
92 5
93 >>> pe.encipher_letter('k')
94 'h'
95 """
96 self.advance()
97 return self.lookup(letter)
98 decipher_letter = encipher_letter
99
100 def lookup(self, letter):
101 """Look up what a letter enciphers to, without turning the wheel.
102
103 >>> pe.set_position('f')
104 5
105 >>> cat([pe.lookup(l) for l in string.ascii_lowercase])
106 'udhbfejcpgmokrliwntsayqzvx'
107 >>> pe.lookup('A')
108 ''
109 """
110 if letter in string.ascii_lowercase:
111 return unpos(
112 (self.wheel_map[(pos(letter) - self.position) % 26] +
113 self.position))
114 else:
115 return ''
116
117 def advance(self):
118 """Advances the wheel one position.
119
120 >>> pe.set_position('f')
121 5
122 >>> pe.advance()
123 6
124 """
125 self.position = (self.position + 1) % 26
126 return self.position
127
128 def encipher(self, message, starting_position=None):
129 """Enciphers a whole message.
130
131 >>> pe.set_position('f')
132 5
133 >>> pe.encipher('helloworld')
134 'kjsglcjoqc'
135 >>> pe.set_position('f')
136 5
137 >>> pe.encipher('kjsglcjoqc')
138 'helloworld'
139 >>> pe.encipher('helloworld', starting_position = 'x')
140 'egrekthnnf'
141 """
142 if starting_position:
143 self.set_position(starting_position)
144 transformed = ''
145 for l in message:
146 transformed += self.encipher_letter(l)
147 return transformed
148 decipher = encipher
149
150 def set_position(self, position):
151 """Sets the position of the wheel, by specifying the letter the arrow
152 points to.
153
154 >>> pe.set_position('a')
155 0
156 >>> pe.set_position('m')
157 12
158 >>> pe.set_position('z')
159 25
160 """
161 self.position = pos(position)
162 return self.position
163
164
165 def pocket_enigma_break_by_crib(message, wheel_spec, crib, crib_position):
166 """Break a pocket enigma using a crib (some plaintext that's expected to
167 be in a certain position). Returns a list of possible starting wheel
168 positions that could produce the crib.
169
170 >>> pocket_enigma_break_by_crib('kzpjlzmoga', 1, 'h', 0)
171 ['a', 'f', 'q']
172 >>> pocket_enigma_break_by_crib('kzpjlzmoga', 1, 'he', 0)
173 ['a']
174 >>> pocket_enigma_break_by_crib('kzpjlzmoga', 1, 'll', 2)
175 ['a']
176 >>> pocket_enigma_break_by_crib('kzpjlzmoga', 1, 'l', 2)
177 ['a']
178 >>> pocket_enigma_break_by_crib('kzpjlzmoga', 1, 'l', 3)
179 ['a', 'j', 'n']
180 >>> pocket_enigma_break_by_crib('aaaaa', 1, 'l', 3)
181 []
182 """
183 pe = PocketEnigma(wheel=wheel_spec)
184 possible_positions = []
185 for p in string.ascii_lowercase:
186 pe.set_position(p)
187 plaintext = pe.decipher(message)
188 if plaintext[crib_position:crib_position+len(crib)] == crib:
189 possible_positions += [p]
190 return possible_positions
191
192 if __name__ == "__main__":
193 import doctest
194 doctest.testmod(extraglobs={'pe': PocketEnigma(1, 'a')})