1656720d64867dc71f547f746f202776792bf633
[cipher-training.git] / cipher.py
1 """A set of ciphers with implementations for both enciphering and deciphering
2 them. See cipherbreak for automatic breaking of these ciphers
3 """
4
5 import string
6 import collections
7 from enum import Enum
8 from itertools import zip_longest, cycle, chain
9 from language_models import unaccent, sanitise
10
11
12 modular_division_table = [[0]*26 for _ in range(26)]
13 for a in range(26):
14 for b in range(26):
15 c = (a * b) % 26
16 modular_division_table[b][c] = a
17
18
19 def every_nth(text, n, fillvalue=''):
20 """Returns n strings, each of which consists of every nth character,
21 starting with the 0th, 1st, 2nd, ... (n-1)th character
22
23 >>> every_nth(string.ascii_lowercase, 5)
24 ['afkpuz', 'bglqv', 'chmrw', 'dinsx', 'ejoty']
25 >>> every_nth(string.ascii_lowercase, 1)
26 ['abcdefghijklmnopqrstuvwxyz']
27 >>> every_nth(string.ascii_lowercase, 26) # doctest: +NORMALIZE_WHITESPACE
28 ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
29 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
30 >>> every_nth(string.ascii_lowercase, 5, fillvalue='!')
31 ['afkpuz', 'bglqv!', 'chmrw!', 'dinsx!', 'ejoty!']
32 """
33 split_text = chunks(text, n, fillvalue)
34 return [''.join(l) for l in zip_longest(*split_text, fillvalue=fillvalue)]
35
36 def combine_every_nth(split_text):
37 """Reforms a text split into every_nth strings
38
39 >>> combine_every_nth(every_nth(string.ascii_lowercase, 5))
40 'abcdefghijklmnopqrstuvwxyz'
41 >>> combine_every_nth(every_nth(string.ascii_lowercase, 1))
42 'abcdefghijklmnopqrstuvwxyz'
43 >>> combine_every_nth(every_nth(string.ascii_lowercase, 26))
44 'abcdefghijklmnopqrstuvwxyz'
45 """
46 return ''.join([''.join(l)
47 for l in zip_longest(*split_text, fillvalue='')])
48
49 def chunks(text, n, fillvalue=None):
50 """Split a text into chunks of n characters
51
52 >>> chunks('abcdefghi', 3)
53 ['abc', 'def', 'ghi']
54 >>> chunks('abcdefghi', 4)
55 ['abcd', 'efgh', 'i']
56 >>> chunks('abcdefghi', 4, fillvalue='!')
57 ['abcd', 'efgh', 'i!!!']
58 """
59 if fillvalue:
60 padding = fillvalue[0] * (n - len(text) % n)
61 else:
62 padding = ''
63 return [(text+padding)[i:i+n] for i in range(0, len(text), n)]
64
65 def transpose(items, transposition):
66 """Moves items around according to the given transposition
67
68 >>> transpose(['a', 'b', 'c', 'd'], (0,1,2,3))
69 ['a', 'b', 'c', 'd']
70 >>> transpose(['a', 'b', 'c', 'd'], (3,1,2,0))
71 ['d', 'b', 'c', 'a']
72 >>> transpose([10,11,12,13,14,15], (3,2,4,1,5,0))
73 [13, 12, 14, 11, 15, 10]
74 """
75 transposed = [''] * len(transposition)
76 for p, t in enumerate(transposition):
77 transposed[p] = items[t]
78 return transposed
79
80 def untranspose(items, transposition):
81 """Undoes a transpose
82
83 >>> untranspose(['a', 'b', 'c', 'd'], [0,1,2,3])
84 ['a', 'b', 'c', 'd']
85 >>> untranspose(['d', 'b', 'c', 'a'], [3,1,2,0])
86 ['a', 'b', 'c', 'd']
87 >>> untranspose([13, 12, 14, 11, 15, 10], [3,2,4,1,5,0])
88 [10, 11, 12, 13, 14, 15]
89 """
90 transposed = [''] * len(transposition)
91 for p, t in enumerate(transposition):
92 transposed[t] = items[p]
93 return transposed
94
95 def deduplicate(text):
96 """If a string contains duplicate letters, remove all but the first. Retain
97 the order of the letters.
98
99 >>> deduplicate('cat')
100 ['c', 'a', 't']
101 >>> deduplicate('happy')
102 ['h', 'a', 'p', 'y']
103 >>> deduplicate('cattca')
104 ['c', 'a', 't']
105 """
106 return list(collections.OrderedDict.fromkeys(text))
107
108
109 def caesar_encipher_letter(accented_letter, shift):
110 """Encipher a letter, given a shift amount
111
112 >>> caesar_encipher_letter('a', 1)
113 'b'
114 >>> caesar_encipher_letter('a', 2)
115 'c'
116 >>> caesar_encipher_letter('b', 2)
117 'd'
118 >>> caesar_encipher_letter('x', 2)
119 'z'
120 >>> caesar_encipher_letter('y', 2)
121 'a'
122 >>> caesar_encipher_letter('z', 2)
123 'b'
124 >>> caesar_encipher_letter('z', -1)
125 'y'
126 >>> caesar_encipher_letter('a', -1)
127 'z'
128 >>> caesar_encipher_letter('A', 1)
129 'B'
130 >>> caesar_encipher_letter('é', 1)
131 'f'
132 """
133 letter = unaccent(accented_letter)
134 if letter in string.ascii_letters:
135 if letter in string.ascii_uppercase:
136 alphabet_start = ord('A')
137 else:
138 alphabet_start = ord('a')
139 return chr(((ord(letter) - alphabet_start + shift) % 26) +
140 alphabet_start)
141 else:
142 return letter
143
144 def caesar_decipher_letter(letter, shift):
145 """Decipher a letter, given a shift amount
146
147 >>> caesar_decipher_letter('b', 1)
148 'a'
149 >>> caesar_decipher_letter('b', 2)
150 'z'
151 """
152 return caesar_encipher_letter(letter, -shift)
153
154 def caesar_encipher(message, shift):
155 """Encipher a message with the Caesar cipher of given shift
156
157 >>> caesar_encipher('abc', 1)
158 'bcd'
159 >>> caesar_encipher('abc', 2)
160 'cde'
161 >>> caesar_encipher('abcxyz', 2)
162 'cdezab'
163 >>> caesar_encipher('ab cx yz', 2)
164 'cd ez ab'
165 >>> caesar_encipher('Héllo World!', 2)
166 'Jgnnq Yqtnf!'
167 """
168 enciphered = [caesar_encipher_letter(l, shift) for l in message]
169 return ''.join(enciphered)
170
171 def caesar_decipher(message, shift):
172 """Decipher a message with the Caesar cipher of given shift
173
174 >>> caesar_decipher('bcd', 1)
175 'abc'
176 >>> caesar_decipher('cde', 2)
177 'abc'
178 >>> caesar_decipher('cd ez ab', 2)
179 'ab cx yz'
180 >>> caesar_decipher('Jgnnq Yqtnf!', 2)
181 'Hello World!'
182 """
183 return caesar_encipher(message, -shift)
184
185 def affine_encipher_letter(accented_letter, multiplier=1, adder=0,
186 one_based=True):
187 """Encipher a letter, given a multiplier and adder
188 >>> ''.join([affine_encipher_letter(l, 3, 5, True) \
189 for l in string.ascii_uppercase])
190 'HKNQTWZCFILORUXADGJMPSVYBE'
191 >>> ''.join([affine_encipher_letter(l, 3, 5, False) \
192 for l in string.ascii_uppercase])
193 'FILORUXADGJMPSVYBEHKNQTWZC'
194 """
195 letter = unaccent(accented_letter)
196 if letter in string.ascii_letters:
197 if letter in string.ascii_uppercase:
198 alphabet_start = ord('A')
199 else:
200 alphabet_start = ord('a')
201 letter_number = ord(letter) - alphabet_start
202 if one_based: letter_number += 1
203 cipher_number = (letter_number * multiplier + adder) % 26
204 if one_based: cipher_number -= 1
205 return chr(cipher_number % 26 + alphabet_start)
206 else:
207 return letter
208
209 def affine_decipher_letter(letter, multiplier=1, adder=0, one_based=True):
210 """Encipher a letter, given a multiplier and adder
211
212 >>> ''.join([affine_decipher_letter(l, 3, 5, True) \
213 for l in 'HKNQTWZCFILORUXADGJMPSVYBE'])
214 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
215 >>> ''.join([affine_decipher_letter(l, 3, 5, False) \
216 for l in 'FILORUXADGJMPSVYBEHKNQTWZC'])
217 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
218 """
219 if letter in string.ascii_letters:
220 if letter in string.ascii_uppercase:
221 alphabet_start = ord('A')
222 else:
223 alphabet_start = ord('a')
224 cipher_number = ord(letter) - alphabet_start
225 if one_based: cipher_number += 1
226 plaintext_number = (
227 modular_division_table[multiplier]
228 [(cipher_number - adder) % 26] )
229 if one_based: plaintext_number -= 1
230 return chr(plaintext_number % 26 + alphabet_start)
231 else:
232 return letter
233
234 def affine_encipher(message, multiplier=1, adder=0, one_based=True):
235 """Encipher a message
236
237 >>> affine_encipher('hours passed during which jerico tried every ' \
238 'trick he could think of', 15, 22, True)
239 'lmyfu bkuusd dyfaxw claol psfaom jfasd snsfg jfaoe ls omytd jlaxe mh'
240 """
241 enciphered = [affine_encipher_letter(l, multiplier, adder, one_based)
242 for l in message]
243 return ''.join(enciphered)
244
245 def affine_decipher(message, multiplier=1, adder=0, one_based=True):
246 """Decipher a message
247
248 >>> affine_decipher('lmyfu bkuusd dyfaxw claol psfaom jfasd snsfg ' \
249 'jfaoe ls omytd jlaxe mh', 15, 22, True)
250 'hours passed during which jerico tried every trick he could think of'
251 """
252 enciphered = [affine_decipher_letter(l, multiplier, adder, one_based)
253 for l in message]
254 return ''.join(enciphered)
255
256
257 class Keyword_wrap_alphabet(Enum):
258 """Ways of wrapping the alphabet for keyword-based substitution ciphers."""
259 from_a = 1
260 from_last = 2
261 from_largest = 3
262
263
264 def keyword_cipher_alphabet_of(keyword,
265 wrap_alphabet=Keyword_wrap_alphabet.from_a):
266 """Find the cipher alphabet given a keyword.
267 wrap_alphabet controls how the rest of the alphabet is added
268 after the keyword.
269
270 >>> keyword_cipher_alphabet_of('bayes')
271 'bayescdfghijklmnopqrtuvwxz'
272 >>> keyword_cipher_alphabet_of('bayes', Keyword_wrap_alphabet.from_a)
273 'bayescdfghijklmnopqrtuvwxz'
274 >>> keyword_cipher_alphabet_of('bayes', Keyword_wrap_alphabet.from_last)
275 'bayestuvwxzcdfghijklmnopqr'
276 >>> keyword_cipher_alphabet_of('bayes', Keyword_wrap_alphabet.from_largest)
277 'bayeszcdfghijklmnopqrtuvwx'
278 """
279 if wrap_alphabet == Keyword_wrap_alphabet.from_a:
280 cipher_alphabet = ''.join(deduplicate(sanitise(keyword) +
281 string.ascii_lowercase))
282 else:
283 if wrap_alphabet == Keyword_wrap_alphabet.from_last:
284 last_keyword_letter = deduplicate(sanitise(keyword))[-1]
285 else:
286 last_keyword_letter = sorted(sanitise(keyword))[-1]
287 last_keyword_position = string.ascii_lowercase.find(
288 last_keyword_letter) + 1
289 cipher_alphabet = ''.join(
290 deduplicate(sanitise(keyword) +
291 string.ascii_lowercase[last_keyword_position:] +
292 string.ascii_lowercase))
293 return cipher_alphabet
294
295
296 def keyword_encipher(message, keyword,
297 wrap_alphabet=Keyword_wrap_alphabet.from_a):
298 """Enciphers a message with a keyword substitution cipher.
299 wrap_alphabet controls how the rest of the alphabet is added
300 after the keyword.
301 0 : from 'a'
302 1 : from the last letter in the sanitised keyword
303 2 : from the largest letter in the sanitised keyword
304
305 >>> keyword_encipher('test message', 'bayes')
306 'rsqr ksqqbds'
307 >>> keyword_encipher('test message', 'bayes', Keyword_wrap_alphabet.from_a)
308 'rsqr ksqqbds'
309 >>> keyword_encipher('test message', 'bayes', Keyword_wrap_alphabet.from_last)
310 'lskl dskkbus'
311 >>> keyword_encipher('test message', 'bayes', Keyword_wrap_alphabet.from_largest)
312 'qspq jsppbcs'
313 """
314 cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
315 cipher_translation = ''.maketrans(string.ascii_lowercase, cipher_alphabet)
316 return unaccent(message).lower().translate(cipher_translation)
317
318 def keyword_decipher(message, keyword, wrap_alphabet=Keyword_wrap_alphabet.from_a):
319 """Deciphers a message with a keyword substitution cipher.
320 wrap_alphabet controls how the rest of the alphabet is added
321 after the keyword.
322 0 : from 'a'
323 1 : from the last letter in the sanitised keyword
324 2 : from the largest letter in the sanitised keyword
325
326 >>> keyword_decipher('rsqr ksqqbds', 'bayes')
327 'test message'
328 >>> keyword_decipher('rsqr ksqqbds', 'bayes', Keyword_wrap_alphabet.from_a)
329 'test message'
330 >>> keyword_decipher('lskl dskkbus', 'bayes', Keyword_wrap_alphabet.from_last)
331 'test message'
332 >>> keyword_decipher('qspq jsppbcs', 'bayes', Keyword_wrap_alphabet.from_largest)
333 'test message'
334 """
335 cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
336 cipher_translation = ''.maketrans(cipher_alphabet, string.ascii_lowercase)
337 return message.lower().translate(cipher_translation)
338
339
340 def vigenere_encipher(message, keyword):
341 """Vigenere encipher
342
343 >>> vigenere_encipher('hello', 'abc')
344 'hfnlp'
345 """
346 shifts = [ord(l) - ord('a') for l in sanitise(keyword)]
347 pairs = zip(message, cycle(shifts))
348 return ''.join([caesar_encipher_letter(l, k) for l, k in pairs])
349
350 def vigenere_decipher(message, keyword):
351 """Vigenere decipher
352
353 >>> vigenere_decipher('hfnlp', 'abc')
354 'hello'
355 """
356 shifts = [ord(l) - ord('a') for l in sanitise(keyword)]
357 pairs = zip(message, cycle(shifts))
358 return ''.join([caesar_decipher_letter(l, k) for l, k in pairs])
359
360 beaufort_encipher=vigenere_decipher
361 beaufort_decipher=vigenere_encipher
362
363
364 def transpositions_of(keyword):
365 """Finds the transpostions given by a keyword. For instance, the keyword
366 'clever' rearranges to 'celrv', so the first column (0) stays first, the
367 second column (1) moves to third, the third column (2) moves to second,
368 and so on.
369
370 If passed a tuple, assume it's already a transposition and just return it.
371
372 >>> transpositions_of('clever')
373 (0, 2, 1, 4, 3)
374 >>> transpositions_of('fred')
375 (3, 2, 0, 1)
376 >>> transpositions_of((3, 2, 0, 1))
377 (3, 2, 0, 1)
378 """
379 if isinstance(keyword, tuple):
380 return keyword
381 else:
382 key = deduplicate(keyword)
383 transpositions = tuple(key.index(l) for l in sorted(key))
384 return transpositions
385
386 def pad(message_len, group_len, fillvalue):
387 """Returns the padding required to extend a message of message_len to an
388 even multiple of group_len, by adding repreated copies of fillvalue.
389 fillvalue can either be a character or a function that returns a character.
390
391 >>> pad(10, 4, '!')
392 '!!'
393 >>> pad(8, 4, '!')
394 ''
395 >>> pad(16, 4, '!')
396 ''
397 >>> pad(10, 4, lambda: '*')
398 '**'
399 """
400 padding_length = group_len - message_len % group_len
401 if padding_length == group_len: padding_length = 0
402 padding = ''
403 for _ in range(padding_length):
404 if callable(fillvalue):
405 padding += fillvalue()
406 else:
407 padding += fillvalue
408 return padding
409
410 def column_transposition_encipher(message, keyword, fillvalue=' ',
411 fillcolumnwise=False,
412 emptycolumnwise=False):
413 """Enciphers using the column transposition cipher.
414 Message is padded to allow all rows to be the same length.
415
416 >>> column_transposition_encipher('hellothere', 'abcdef', fillcolumnwise=True)
417 'hlohr eltee '
418 >>> column_transposition_encipher('hellothere', 'abcdef', fillcolumnwise=True, emptycolumnwise=True)
419 'hellothere '
420 >>> column_transposition_encipher('hellothere', 'abcdef')
421 'hellothere '
422 >>> column_transposition_encipher('hellothere', 'abcde')
423 'hellothere'
424 >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=True)
425 'hellothere'
426 >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=False)
427 'hlohreltee'
428 >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=True)
429 'htehlelroe'
430 >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=False)
431 'hellothere'
432 >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=True, emptycolumnwise=True)
433 'heotllrehe'
434 >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=True, emptycolumnwise=False)
435 'holrhetlee'
436 >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=False, emptycolumnwise=True)
437 'htleehoelr'
438 >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=False, emptycolumnwise=False)
439 'hleolteher'
440 >>> column_transposition_encipher('hellothere', 'cleverly')
441 'hleolthre e '
442 >>> column_transposition_encipher('hellothere', 'cleverly', fillvalue='!')
443 'hleolthre!e!'
444 >>> column_transposition_encipher('hellothere', 'cleverly', fillvalue=lambda: '*')
445 'hleolthre*e*'
446 """
447 transpositions = transpositions_of(keyword)
448 message += pad(len(message), len(transpositions), fillvalue)
449 if fillcolumnwise:
450 rows = every_nth(message, len(message) // len(transpositions))
451 else:
452 rows = chunks(message, len(transpositions))
453 transposed = [transpose(r, transpositions) for r in rows]
454 if emptycolumnwise:
455 return combine_every_nth(transposed)
456 else:
457 return ''.join(chain(*transposed))
458
459 def column_transposition_decipher(message, keyword, fillvalue=' ',
460 fillcolumnwise=False,
461 emptycolumnwise=False):
462 """Deciphers using the column transposition cipher.
463 Message is padded to allow all rows to be the same length.
464
465 >>> column_transposition_decipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=True)
466 'hellothere'
467 >>> column_transposition_decipher('hlohreltee', 'abcde', fillcolumnwise=True, emptycolumnwise=False)
468 'hellothere'
469 >>> column_transposition_decipher('htehlelroe', 'abcde', fillcolumnwise=False, emptycolumnwise=True)
470 'hellothere'
471 >>> column_transposition_decipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=False)
472 'hellothere'
473 >>> column_transposition_decipher('heotllrehe', 'clever', fillcolumnwise=True, emptycolumnwise=True)
474 'hellothere'
475 >>> column_transposition_decipher('holrhetlee', 'clever', fillcolumnwise=True, emptycolumnwise=False)
476 'hellothere'
477 >>> column_transposition_decipher('htleehoelr', 'clever', fillcolumnwise=False, emptycolumnwise=True)
478 'hellothere'
479 >>> column_transposition_decipher('hleolteher', 'clever', fillcolumnwise=False, emptycolumnwise=False)
480 'hellothere'
481 """
482 transpositions = transpositions_of(keyword)
483 message += pad(len(message), len(transpositions), '*')
484 if emptycolumnwise:
485 rows = every_nth(message, len(message) // len(transpositions))
486 else:
487 rows = chunks(message, len(transpositions))
488 untransposed = [untranspose(r, transpositions) for r in rows]
489 if fillcolumnwise:
490 return combine_every_nth(untransposed)
491 else:
492 return ''.join(chain(*untransposed))
493
494 def scytale_encipher(message, rows, fillvalue=' '):
495 """Enciphers using the scytale transposition cipher.
496 Message is padded with spaces to allow all rows to be the same length.
497
498 >>> scytale_encipher('thequickbrownfox', 3)
499 'tcnhkfeboqrxuo iw '
500 >>> scytale_encipher('thequickbrownfox', 4)
501 'tubnhirfecooqkwx'
502 >>> scytale_encipher('thequickbrownfox', 5)
503 'tubn hirf ecoo qkwx '
504 >>> scytale_encipher('thequickbrownfox', 6)
505 'tqcrnxhukof eibwo '
506 >>> scytale_encipher('thequickbrownfox', 7)
507 'tqcrnx hukof eibwo '
508 """
509 transpositions = [i for i in range(rows)]
510 return column_transposition_encipher(message, transpositions,
511 fillvalue=fillvalue, fillcolumnwise=True, emptycolumnwise=False)
512
513 def scytale_decipher(message, rows):
514 """Deciphers using the scytale transposition cipher.
515 Assumes the message is padded so that all rows are the same length.
516
517 >>> scytale_decipher('tcnhkfeboqrxuo iw ', 3)
518 'thequickbrownfox '
519 >>> scytale_decipher('tubnhirfecooqkwx', 4)
520 'thequickbrownfox'
521 >>> scytale_decipher('tubn hirf ecoo qkwx ', 5)
522 'thequickbrownfox '
523 >>> scytale_decipher('tqcrnxhukof eibwo ', 6)
524 'thequickbrownfox '
525 >>> scytale_decipher('tqcrnx hukof eibwo ', 7)
526 'thequickbrownfox '
527 """
528 transpositions = [i for i in range(rows)]
529 return column_transposition_decipher(message, transpositions,
530 fillcolumnwise=True, emptycolumnwise=False)
531
532
533 class PocketEnigma(object):
534 """A pocket enigma machine
535 The wheel is internally represented as a 26-element list self.wheel_map,
536 where wheel_map[i] == j shows that the position i places on from the arrow
537 maps to the position j places on.
538 """
539 def __init__(self, wheel=1, position='a'):
540 """initialise the pocket enigma, including which wheel to use and the
541 starting position of the wheel.
542
543 The wheel is either 1 or 2 (the predefined wheels) or a list of letter
544 pairs.
545
546 The position is the letter pointed to by the arrow on the wheel.
547
548 >>> pe.wheel_map
549 [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]
550 >>> pe.position
551 0
552 """
553 self.wheel1 = [('a', 'z'), ('b', 'e'), ('c', 'x'), ('d', 'k'),
554 ('f', 'h'), ('g', 'j'), ('i', 'm'), ('l', 'r'), ('n', 'o'),
555 ('p', 'v'), ('q', 't'), ('s', 'u'), ('w', 'y')]
556 self.wheel2 = [('a', 'c'), ('b', 'd'), ('e', 'w'), ('f', 'i'),
557 ('g', 'p'), ('h', 'm'), ('j', 'k'), ('l', 'n'), ('o', 'q'),
558 ('r', 'z'), ('s', 'u'), ('t', 'v'), ('x', 'y')]
559 if wheel == 1:
560 self.make_wheel_map(self.wheel1)
561 elif wheel == 2:
562 self.make_wheel_map(self.wheel2)
563 else:
564 self.validate_wheel_spec(wheel)
565 self.make_wheel_map(wheel)
566 if position in string.ascii_lowercase:
567 self.position = ord(position) - ord('a')
568 else:
569 self.position = position
570
571 def make_wheel_map(self, wheel_spec):
572 """Expands a wheel specification from a list of letter-letter pairs
573 into a full wheel_map.
574
575 >>> pe.make_wheel_map(pe.wheel2)
576 [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]
577 """
578 self.validate_wheel_spec(wheel_spec)
579 self.wheel_map = [0] * 26
580 for p in wheel_spec:
581 self.wheel_map[ord(p[0]) - ord('a')] = ord(p[1]) - ord('a')
582 self.wheel_map[ord(p[1]) - ord('a')] = ord(p[0]) - ord('a')
583 return self.wheel_map
584
585 def validate_wheel_spec(self, wheel_spec):
586 """Validates that a wheel specificaiton will turn into a valid wheel
587 map.
588
589 >>> pe.validate_wheel_spec([])
590 Traceback (most recent call last):
591 ...
592 ValueError: Wheel specification has 0 pairs, requires 13
593 >>> pe.validate_wheel_spec([('a', 'b', 'c')]*13)
594 Traceback (most recent call last):
595 ...
596 ValueError: Not all mappings in wheel specificationhave two elements
597 >>> pe.validate_wheel_spec([('a', 'b')]*13)
598 Traceback (most recent call last):
599 ...
600 ValueError: Wheel specification does not contain 26 letters
601 """
602 if len(wheel_spec) != 13:
603 raise ValueError("Wheel specification has {} pairs, requires 13".
604 format(len(wheel_spec)))
605 for p in wheel_spec:
606 if len(p) != 2:
607 raise ValueError("Not all mappings in wheel specification"
608 "have two elements")
609 if len(set([p[0] for p in wheel_spec] +
610 [p[1] for p in wheel_spec])) != 26:
611 raise ValueError("Wheel specification does not contain 26 letters")
612
613 def encipher_letter(self, letter):
614 """Enciphers a single letter, by advancing the wheel before looking up
615 the letter on the wheel.
616
617 >>> pe.set_position('f')
618 5
619 >>> pe.encipher_letter('k')
620 'h'
621 """
622 self.advance()
623 return self.lookup(letter)
624 decipher_letter = encipher_letter
625
626 def lookup(self, letter):
627 """Look up what a letter enciphers to, without turning the wheel.
628
629 >>> pe.set_position('f')
630 5
631 >>> ''.join([pe.lookup(l) for l in string.ascii_lowercase])
632 'udhbfejcpgmokrliwntsayqzvx'
633 >>> pe.lookup('A')
634 ''
635 """
636 if letter in string.ascii_lowercase:
637 return chr(
638 (self.wheel_map[(ord(letter) - ord('a') - self.position) % 26] +
639 self.position) % 26 +
640 ord('a'))
641 else:
642 return ''
643
644 def advance(self):
645 """Advances the wheel one position.
646
647 >>> pe.set_position('f')
648 5
649 >>> pe.advance()
650 6
651 """
652 self.position = (self.position + 1) % 26
653 return self.position
654
655 def encipher(self, message, starting_position=None):
656 """Enciphers a whole message.
657
658 >>> pe.set_position('f')
659 5
660 >>> pe.encipher('helloworld')
661 'kjsglcjoqc'
662 >>> pe.set_position('f')
663 5
664 >>> pe.encipher('kjsglcjoqc')
665 'helloworld'
666 >>> pe.encipher('helloworld', starting_position = 'x')
667 'egrekthnnf'
668 """
669 if starting_position:
670 self.set_position(starting_position)
671 transformed = ''
672 for l in message:
673 transformed += self.encipher_letter(l)
674 return transformed
675 decipher = encipher
676
677 def set_position(self, position):
678 """Sets the position of the wheel, by specifying the letter the arrow
679 points to.
680
681 >>> pe.set_position('a')
682 0
683 >>> pe.set_position('m')
684 12
685 >>> pe.set_position('z')
686 25
687 """
688 self.position = ord(position) - ord('a')
689 return self.position
690
691
692 if __name__ == "__main__":
693 import doctest
694 doctest.testmod(extraglobs={'pe': PocketEnigma(1, 'a')})