More tweaking to conform with linting
[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 )
230 if one_based: plaintext_number -= 1
231 return chr(plaintext_number % 26 + alphabet_start)
232 else:
233 return letter
234
235 def affine_encipher(message, multiplier=1, adder=0, one_based=True):
236 """Encipher a message
237
238 >>> affine_encipher('hours passed during which jerico tried every ' \
239 'trick he could think of', 15, 22, True)
240 'lmyfu bkuusd dyfaxw claol psfaom jfasd snsfg jfaoe ls omytd jlaxe mh'
241 """
242 enciphered = [affine_encipher_letter(l, multiplier, adder, one_based)
243 for l in message]
244 return ''.join(enciphered)
245
246 def affine_decipher(message, multiplier=1, adder=0, one_based=True):
247 """Decipher a message
248
249 >>> affine_decipher('lmyfu bkuusd dyfaxw claol psfaom jfasd snsfg ' \
250 'jfaoe ls omytd jlaxe mh', 15, 22, True)
251 'hours passed during which jerico tried every trick he could think of'
252 """
253 enciphered = [affine_decipher_letter(l, multiplier, adder, one_based)
254 for l in message]
255 return ''.join(enciphered)
256
257
258 class KeywordWrapAlphabet(Enum):
259 """Ways of wrapping the alphabet for keyword-based substitution ciphers."""
260 from_a = 1
261 from_last = 2
262 from_largest = 3
263
264
265 def keyword_cipher_alphabet_of(keyword,
266 wrap_alphabet=KeywordWrapAlphabet.from_a):
267 """Find the cipher alphabet given a keyword.
268 wrap_alphabet controls how the rest of the alphabet is added
269 after the keyword.
270
271 >>> keyword_cipher_alphabet_of('bayes')
272 'bayescdfghijklmnopqrtuvwxz'
273 >>> keyword_cipher_alphabet_of('bayes', KeywordWrapAlphabet.from_a)
274 'bayescdfghijklmnopqrtuvwxz'
275 >>> keyword_cipher_alphabet_of('bayes', KeywordWrapAlphabet.from_last)
276 'bayestuvwxzcdfghijklmnopqr'
277 >>> keyword_cipher_alphabet_of('bayes', KeywordWrapAlphabet.from_largest)
278 'bayeszcdfghijklmnopqrtuvwx'
279 """
280 if wrap_alphabet == KeywordWrapAlphabet.from_a:
281 cipher_alphabet = ''.join(deduplicate(sanitise(keyword) +
282 string.ascii_lowercase))
283 else:
284 if wrap_alphabet == KeywordWrapAlphabet.from_last:
285 last_keyword_letter = deduplicate(sanitise(keyword))[-1]
286 else:
287 last_keyword_letter = sorted(sanitise(keyword))[-1]
288 last_keyword_position = string.ascii_lowercase.find(
289 last_keyword_letter) + 1
290 cipher_alphabet = ''.join(
291 deduplicate(sanitise(keyword) +
292 string.ascii_lowercase[last_keyword_position:] +
293 string.ascii_lowercase))
294 return cipher_alphabet
295
296
297 def keyword_encipher(message, keyword,
298 wrap_alphabet=KeywordWrapAlphabet.from_a):
299 """Enciphers a message with a keyword substitution cipher.
300 wrap_alphabet controls how the rest of the alphabet is added
301 after the keyword.
302 0 : from 'a'
303 1 : from the last letter in the sanitised keyword
304 2 : from the largest letter in the sanitised keyword
305
306 >>> keyword_encipher('test message', 'bayes')
307 'rsqr ksqqbds'
308 >>> keyword_encipher('test message', 'bayes', KeywordWrapAlphabet.from_a)
309 'rsqr ksqqbds'
310 >>> keyword_encipher('test message', 'bayes', KeywordWrapAlphabet.from_last)
311 'lskl dskkbus'
312 >>> keyword_encipher('test message', 'bayes', KeywordWrapAlphabet.from_largest)
313 'qspq jsppbcs'
314 """
315 cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
316 cipher_translation = ''.maketrans(string.ascii_lowercase, cipher_alphabet)
317 return unaccent(message).lower().translate(cipher_translation)
318
319 def keyword_decipher(message, keyword,
320 wrap_alphabet=KeywordWrapAlphabet.from_a):
321 """Deciphers a message with a keyword substitution cipher.
322 wrap_alphabet controls how the rest of the alphabet is added
323 after the keyword.
324 0 : from 'a'
325 1 : from the last letter in the sanitised keyword
326 2 : from the largest letter in the sanitised keyword
327
328 >>> keyword_decipher('rsqr ksqqbds', 'bayes')
329 'test message'
330 >>> keyword_decipher('rsqr ksqqbds', 'bayes', KeywordWrapAlphabet.from_a)
331 'test message'
332 >>> keyword_decipher('lskl dskkbus', 'bayes', KeywordWrapAlphabet.from_last)
333 'test message'
334 >>> keyword_decipher('qspq jsppbcs', 'bayes', KeywordWrapAlphabet.from_largest)
335 'test message'
336 """
337 cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
338 cipher_translation = ''.maketrans(cipher_alphabet, string.ascii_lowercase)
339 return message.lower().translate(cipher_translation)
340
341
342 def vigenere_encipher(message, keyword):
343 """Vigenere encipher
344
345 >>> vigenere_encipher('hello', 'abc')
346 'hfnlp'
347 """
348 shifts = [ord(l) - ord('a') for l in sanitise(keyword)]
349 pairs = zip(message, cycle(shifts))
350 return ''.join([caesar_encipher_letter(l, k) for l, k in pairs])
351
352 def vigenere_decipher(message, keyword):
353 """Vigenere decipher
354
355 >>> vigenere_decipher('hfnlp', 'abc')
356 'hello'
357 """
358 shifts = [ord(l) - ord('a') for l in sanitise(keyword)]
359 pairs = zip(message, cycle(shifts))
360 return ''.join([caesar_decipher_letter(l, k) for l, k in pairs])
361
362 beaufort_encipher = vigenere_decipher
363 beaufort_decipher = vigenere_encipher
364
365
366 def transpositions_of(keyword):
367 """Finds the transpostions given by a keyword. For instance, the keyword
368 'clever' rearranges to 'celrv', so the first column (0) stays first, the
369 second column (1) moves to third, the third column (2) moves to second,
370 and so on.
371
372 If passed a tuple, assume it's already a transposition and just return it.
373
374 >>> transpositions_of('clever')
375 (0, 2, 1, 4, 3)
376 >>> transpositions_of('fred')
377 (3, 2, 0, 1)
378 >>> transpositions_of((3, 2, 0, 1))
379 (3, 2, 0, 1)
380 """
381 if isinstance(keyword, tuple):
382 return keyword
383 else:
384 key = deduplicate(keyword)
385 transpositions = tuple(key.index(l) for l in sorted(key))
386 return transpositions
387
388 def pad(message_len, group_len, fillvalue):
389 """Returns the padding required to extend a message of message_len to an
390 even multiple of group_len, by adding repreated copies of fillvalue.
391 fillvalue can either be a character or a function that returns a character.
392
393 >>> pad(10, 4, '!')
394 '!!'
395 >>> pad(8, 4, '!')
396 ''
397 >>> pad(16, 4, '!')
398 ''
399 >>> pad(10, 4, lambda: '*')
400 '**'
401 """
402 padding_length = group_len - message_len % group_len
403 if padding_length == group_len: padding_length = 0
404 padding = ''
405 for _ in range(padding_length):
406 if callable(fillvalue):
407 padding += fillvalue()
408 else:
409 padding += fillvalue
410 return padding
411
412 def column_transposition_encipher(message, keyword, fillvalue=' ',
413 fillcolumnwise=False,
414 emptycolumnwise=False):
415 """Enciphers using the column transposition cipher.
416 Message is padded to allow all rows to be the same length.
417
418 >>> column_transposition_encipher('hellothere', 'abcdef', fillcolumnwise=True)
419 'hlohr eltee '
420 >>> column_transposition_encipher('hellothere', 'abcdef', fillcolumnwise=True, emptycolumnwise=True)
421 'hellothere '
422 >>> column_transposition_encipher('hellothere', 'abcdef')
423 'hellothere '
424 >>> column_transposition_encipher('hellothere', 'abcde')
425 'hellothere'
426 >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=True)
427 'hellothere'
428 >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=False)
429 'hlohreltee'
430 >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=True)
431 'htehlelroe'
432 >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=False)
433 'hellothere'
434 >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=True, emptycolumnwise=True)
435 'heotllrehe'
436 >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=True, emptycolumnwise=False)
437 'holrhetlee'
438 >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=False, emptycolumnwise=True)
439 'htleehoelr'
440 >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=False, emptycolumnwise=False)
441 'hleolteher'
442 >>> column_transposition_encipher('hellothere', 'cleverly')
443 'hleolthre e '
444 >>> column_transposition_encipher('hellothere', 'cleverly', fillvalue='!')
445 'hleolthre!e!'
446 >>> column_transposition_encipher('hellothere', 'cleverly', fillvalue=lambda: '*')
447 'hleolthre*e*'
448 """
449 transpositions = transpositions_of(keyword)
450 message += pad(len(message), len(transpositions), fillvalue)
451 if fillcolumnwise:
452 rows = every_nth(message, len(message) // len(transpositions))
453 else:
454 rows = chunks(message, len(transpositions))
455 transposed = [transpose(r, transpositions) for r in rows]
456 if emptycolumnwise:
457 return combine_every_nth(transposed)
458 else:
459 return ''.join(chain(*transposed))
460
461 def column_transposition_decipher(message, keyword, fillvalue=' ',
462 fillcolumnwise=False,
463 emptycolumnwise=False):
464 """Deciphers using the column transposition cipher.
465 Message is padded to allow all rows to be the same length.
466
467 >>> column_transposition_decipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=True)
468 'hellothere'
469 >>> column_transposition_decipher('hlohreltee', 'abcde', fillcolumnwise=True, emptycolumnwise=False)
470 'hellothere'
471 >>> column_transposition_decipher('htehlelroe', 'abcde', fillcolumnwise=False, emptycolumnwise=True)
472 'hellothere'
473 >>> column_transposition_decipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=False)
474 'hellothere'
475 >>> column_transposition_decipher('heotllrehe', 'clever', fillcolumnwise=True, emptycolumnwise=True)
476 'hellothere'
477 >>> column_transposition_decipher('holrhetlee', 'clever', fillcolumnwise=True, emptycolumnwise=False)
478 'hellothere'
479 >>> column_transposition_decipher('htleehoelr', 'clever', fillcolumnwise=False, emptycolumnwise=True)
480 'hellothere'
481 >>> column_transposition_decipher('hleolteher', 'clever', fillcolumnwise=False, emptycolumnwise=False)
482 'hellothere'
483 """
484 transpositions = transpositions_of(keyword)
485 message += pad(len(message), len(transpositions), '*')
486 if emptycolumnwise:
487 rows = every_nth(message, len(message) // len(transpositions))
488 else:
489 rows = chunks(message, len(transpositions))
490 untransposed = [untranspose(r, transpositions) for r in rows]
491 if fillcolumnwise:
492 return combine_every_nth(untransposed)
493 else:
494 return ''.join(chain(*untransposed))
495
496 def scytale_encipher(message, rows, fillvalue=' '):
497 """Enciphers using the scytale transposition cipher.
498 Message is padded with spaces to allow all rows to be the same length.
499
500 >>> scytale_encipher('thequickbrownfox', 3)
501 'tcnhkfeboqrxuo iw '
502 >>> scytale_encipher('thequickbrownfox', 4)
503 'tubnhirfecooqkwx'
504 >>> scytale_encipher('thequickbrownfox', 5)
505 'tubn hirf ecoo qkwx '
506 >>> scytale_encipher('thequickbrownfox', 6)
507 'tqcrnxhukof eibwo '
508 >>> scytale_encipher('thequickbrownfox', 7)
509 'tqcrnx hukof eibwo '
510 """
511 transpositions = [i for i in range(rows)]
512 return column_transposition_encipher(message, transpositions,
513 fillvalue=fillvalue, fillcolumnwise=True, emptycolumnwise=False)
514
515 def scytale_decipher(message, rows):
516 """Deciphers using the scytale transposition cipher.
517 Assumes the message is padded so that all rows are the same length.
518
519 >>> scytale_decipher('tcnhkfeboqrxuo iw ', 3)
520 'thequickbrownfox '
521 >>> scytale_decipher('tubnhirfecooqkwx', 4)
522 'thequickbrownfox'
523 >>> scytale_decipher('tubn hirf ecoo qkwx ', 5)
524 'thequickbrownfox '
525 >>> scytale_decipher('tqcrnxhukof eibwo ', 6)
526 'thequickbrownfox '
527 >>> scytale_decipher('tqcrnx hukof eibwo ', 7)
528 'thequickbrownfox '
529 """
530 transpositions = [i for i in range(rows)]
531 return column_transposition_decipher(message, transpositions,
532 fillcolumnwise=True, emptycolumnwise=False)
533
534
535 class PocketEnigma(object):
536 """A pocket enigma machine
537 The wheel is internally represented as a 26-element list self.wheel_map,
538 where wheel_map[i] == j shows that the position i places on from the arrow
539 maps to the position j places on.
540 """
541 def __init__(self, wheel=1, position='a'):
542 """initialise the pocket enigma, including which wheel to use and the
543 starting position of the wheel.
544
545 The wheel is either 1 or 2 (the predefined wheels) or a list of letter
546 pairs.
547
548 The position is the letter pointed to by the arrow on the wheel.
549
550 >>> pe.wheel_map
551 [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]
552 >>> pe.position
553 0
554 """
555 self.wheel1 = [('a', 'z'), ('b', 'e'), ('c', 'x'), ('d', 'k'),
556 ('f', 'h'), ('g', 'j'), ('i', 'm'), ('l', 'r'), ('n', 'o'),
557 ('p', 'v'), ('q', 't'), ('s', 'u'), ('w', 'y')]
558 self.wheel2 = [('a', 'c'), ('b', 'd'), ('e', 'w'), ('f', 'i'),
559 ('g', 'p'), ('h', 'm'), ('j', 'k'), ('l', 'n'), ('o', 'q'),
560 ('r', 'z'), ('s', 'u'), ('t', 'v'), ('x', 'y')]
561 if wheel == 1:
562 self.make_wheel_map(self.wheel1)
563 elif wheel == 2:
564 self.make_wheel_map(self.wheel2)
565 else:
566 self.validate_wheel_spec(wheel)
567 self.make_wheel_map(wheel)
568 if position in string.ascii_lowercase:
569 self.position = ord(position) - ord('a')
570 else:
571 self.position = position
572
573 def make_wheel_map(self, wheel_spec):
574 """Expands a wheel specification from a list of letter-letter pairs
575 into a full wheel_map.
576
577 >>> pe.make_wheel_map(pe.wheel2)
578 [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]
579 """
580 self.validate_wheel_spec(wheel_spec)
581 self.wheel_map = [0] * 26
582 for p in wheel_spec:
583 self.wheel_map[ord(p[0]) - ord('a')] = ord(p[1]) - ord('a')
584 self.wheel_map[ord(p[1]) - ord('a')] = ord(p[0]) - ord('a')
585 return self.wheel_map
586
587 def validate_wheel_spec(self, wheel_spec):
588 """Validates that a wheel specificaiton will turn into a valid wheel
589 map.
590
591 >>> pe.validate_wheel_spec([])
592 Traceback (most recent call last):
593 ...
594 ValueError: Wheel specification has 0 pairs, requires 13
595 >>> pe.validate_wheel_spec([('a', 'b', 'c')]*13)
596 Traceback (most recent call last):
597 ...
598 ValueError: Not all mappings in wheel specificationhave two elements
599 >>> pe.validate_wheel_spec([('a', 'b')]*13)
600 Traceback (most recent call last):
601 ...
602 ValueError: Wheel specification does not contain 26 letters
603 """
604 if len(wheel_spec) != 13:
605 raise ValueError("Wheel specification has {} pairs, requires"
606 " 13".format(len(wheel_spec)))
607 for p in wheel_spec:
608 if len(p) != 2:
609 raise ValueError("Not all mappings in wheel specification"
610 "have two elements")
611 if len(set([p[0] for p in wheel_spec] +
612 [p[1] for p in wheel_spec])) != 26:
613 raise ValueError("Wheel specification does not contain 26 letters")
614
615 def encipher_letter(self, letter):
616 """Enciphers a single letter, by advancing the wheel before looking up
617 the letter on the wheel.
618
619 >>> pe.set_position('f')
620 5
621 >>> pe.encipher_letter('k')
622 'h'
623 """
624 self.advance()
625 return self.lookup(letter)
626 decipher_letter = encipher_letter
627
628 def lookup(self, letter):
629 """Look up what a letter enciphers to, without turning the wheel.
630
631 >>> pe.set_position('f')
632 5
633 >>> ''.join([pe.lookup(l) for l in string.ascii_lowercase])
634 'udhbfejcpgmokrliwntsayqzvx'
635 >>> pe.lookup('A')
636 ''
637 """
638 if letter in string.ascii_lowercase:
639 return chr(
640 (self.wheel_map[(ord(letter) - ord('a') - self.position) % 26] +
641 self.position) % 26 +
642 ord('a'))
643 else:
644 return ''
645
646 def advance(self):
647 """Advances the wheel one position.
648
649 >>> pe.set_position('f')
650 5
651 >>> pe.advance()
652 6
653 """
654 self.position = (self.position + 1) % 26
655 return self.position
656
657 def encipher(self, message, starting_position=None):
658 """Enciphers a whole message.
659
660 >>> pe.set_position('f')
661 5
662 >>> pe.encipher('helloworld')
663 'kjsglcjoqc'
664 >>> pe.set_position('f')
665 5
666 >>> pe.encipher('kjsglcjoqc')
667 'helloworld'
668 >>> pe.encipher('helloworld', starting_position = 'x')
669 'egrekthnnf'
670 """
671 if starting_position:
672 self.set_position(starting_position)
673 transformed = ''
674 for l in message:
675 transformed += self.encipher_letter(l)
676 return transformed
677 decipher = encipher
678
679 def set_position(self, position):
680 """Sets the position of the wheel, by specifying the letter the arrow
681 points to.
682
683 >>> pe.set_position('a')
684 0
685 >>> pe.set_position('m')
686 12
687 >>> pe.set_position('z')
688 25
689 """
690 self.position = ord(position) - ord('a')
691 return self.position
692
693
694 if __name__ == "__main__":
695 import doctest
696 doctest.testmod(extraglobs={'pe': PocketEnigma(1, 'a')})