Added another test case
[cipher-training.git] / cipher.py
1 import string
2 import collections
3 import logging
4 import math
5 from itertools import zip_longest, cycle, chain
6 from language_models import *
7
8 logger = logging.getLogger(__name__)
9 logger.addHandler(logging.FileHandler('cipher.log'))
10 logger.setLevel(logging.WARNING)
11 #logger.setLevel(logging.INFO)
12 #logger.setLevel(logging.DEBUG)
13
14
15 modular_division_table = [[0]*26 for _ in range(26)]
16 for a in range(26):
17 for b in range(26):
18 c = (a * b) % 26
19 modular_division_table[b][c] = a
20
21
22 def every_nth(text, n, fillvalue=''):
23 """Returns n strings, each of which consists of every nth character,
24 starting with the 0th, 1st, 2nd, ... (n-1)th character
25
26 >>> every_nth(string.ascii_lowercase, 5)
27 ['afkpuz', 'bglqv', 'chmrw', 'dinsx', 'ejoty']
28 >>> every_nth(string.ascii_lowercase, 1)
29 ['abcdefghijklmnopqrstuvwxyz']
30 >>> every_nth(string.ascii_lowercase, 26) # doctest: +NORMALIZE_WHITESPACE
31 ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
32 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
33 >>> every_nth(string.ascii_lowercase, 5, fillvalue='!')
34 ['afkpuz', 'bglqv!', 'chmrw!', 'dinsx!', 'ejoty!']
35 """
36 split_text = [text[i:i+n] for i in range(0, len(text), n)]
37 return [''.join(l) for l in zip_longest(*split_text, fillvalue=fillvalue)]
38
39 def combine_every_nth(split_text):
40 """Reforms a text split into every_nth strings
41
42 >>> combine_every_nth(every_nth(string.ascii_lowercase, 5))
43 'abcdefghijklmnopqrstuvwxyz'
44 >>> combine_every_nth(every_nth(string.ascii_lowercase, 1))
45 'abcdefghijklmnopqrstuvwxyz'
46 >>> combine_every_nth(every_nth(string.ascii_lowercase, 26))
47 'abcdefghijklmnopqrstuvwxyz'
48 """
49 return ''.join([''.join(l)
50 for l in zip_longest(*split_text, fillvalue='')])
51
52 def chunks(text, n, fillvalue=None):
53 """Split a text into chunks of n characters
54
55 >>> chunks('abcdefghi', 3)
56 ['abc', 'def', 'ghi']
57 >>> chunks('abcdefghi', 4)
58 ['abcd', 'efgh', 'i']
59 >>> chunks('abcdefghi', 4, fillvalue='!')
60 ['abcd', 'efgh', 'i!!!']
61 """
62 if fillvalue:
63 padding = fillvalue[0] * (n - len(text) % n)
64 else:
65 padding = ''
66 return [(text+padding)[i:i+n] for i in range(0, len(text), n)]
67
68 def transpose(items, transposition):
69 """Moves items around according to the given transposition
70
71 >>> transpose(['a', 'b', 'c', 'd'], (0,1,2,3))
72 ['a', 'b', 'c', 'd']
73 >>> transpose(['a', 'b', 'c', 'd'], (3,1,2,0))
74 ['d', 'b', 'c', 'a']
75 >>> transpose([10,11,12,13,14,15], (3,2,4,1,5,0))
76 [13, 12, 14, 11, 15, 10]
77 """
78 transposed = [''] * len(transposition)
79 for p, t in enumerate(transposition):
80 transposed[p] = items[t]
81 return transposed
82
83 def untranspose(items, transposition):
84 """Undoes a transpose
85
86 >>> untranspose(['a', 'b', 'c', 'd'], [0,1,2,3])
87 ['a', 'b', 'c', 'd']
88 >>> untranspose(['d', 'b', 'c', 'a'], [3,1,2,0])
89 ['a', 'b', 'c', 'd']
90 >>> untranspose([13, 12, 14, 11, 15, 10], [3,2,4,1,5,0])
91 [10, 11, 12, 13, 14, 15]
92 """
93 transposed = [''] * len(transposition)
94 for p, t in enumerate(transposition):
95 transposed[t] = items[p]
96 return transposed
97
98 def deduplicate(text):
99 return list(collections.OrderedDict.fromkeys(text))
100
101
102 def caesar_encipher_letter(accented_letter, shift):
103 """Encipher a letter, given a shift amount
104
105 >>> caesar_encipher_letter('a', 1)
106 'b'
107 >>> caesar_encipher_letter('a', 2)
108 'c'
109 >>> caesar_encipher_letter('b', 2)
110 'd'
111 >>> caesar_encipher_letter('x', 2)
112 'z'
113 >>> caesar_encipher_letter('y', 2)
114 'a'
115 >>> caesar_encipher_letter('z', 2)
116 'b'
117 >>> caesar_encipher_letter('z', -1)
118 'y'
119 >>> caesar_encipher_letter('a', -1)
120 'z'
121 >>> caesar_encipher_letter('A', 1)
122 'B'
123 >>> caesar_encipher_letter('é', 1)
124 'f'
125 """
126 letter = unaccent(accented_letter)
127 if letter in string.ascii_letters:
128 if letter in string.ascii_uppercase:
129 alphabet_start = ord('A')
130 else:
131 alphabet_start = ord('a')
132 return chr(((ord(letter) - alphabet_start + shift) % 26) +
133 alphabet_start)
134 else:
135 return letter
136
137 def caesar_decipher_letter(letter, shift):
138 """Decipher a letter, given a shift amount
139
140 >>> caesar_decipher_letter('b', 1)
141 'a'
142 >>> caesar_decipher_letter('b', 2)
143 'z'
144 """
145 return caesar_encipher_letter(letter, -shift)
146
147 def caesar_encipher(message, shift):
148 """Encipher a message with the Caesar cipher of given shift
149
150 >>> caesar_encipher('abc', 1)
151 'bcd'
152 >>> caesar_encipher('abc', 2)
153 'cde'
154 >>> caesar_encipher('abcxyz', 2)
155 'cdezab'
156 >>> caesar_encipher('ab cx yz', 2)
157 'cd ez ab'
158 >>> caesar_encipher('Héllo World!', 2)
159 'Jgnnq Yqtnf!'
160 """
161 enciphered = [caesar_encipher_letter(l, shift) for l in message]
162 return ''.join(enciphered)
163
164 def caesar_decipher(message, shift):
165 """Decipher a message with the Caesar cipher of given shift
166
167 >>> caesar_decipher('bcd', 1)
168 'abc'
169 >>> caesar_decipher('cde', 2)
170 'abc'
171 >>> caesar_decipher('cd ez ab', 2)
172 'ab cx yz'
173 >>> caesar_decipher('Jgnnq Yqtnf!', 2)
174 'Hello World!'
175 """
176 return caesar_encipher(message, -shift)
177
178 def affine_encipher_letter(accented_letter, multiplier=1, adder=0, one_based=True):
179 """Encipher a letter, given a multiplier and adder
180
181 >>> ''.join([affine_encipher_letter(l, 3, 5, True) \
182 for l in string.ascii_uppercase])
183 'HKNQTWZCFILORUXADGJMPSVYBE'
184 >>> ''.join([affine_encipher_letter(l, 3, 5, False) \
185 for l in string.ascii_uppercase])
186 'FILORUXADGJMPSVYBEHKNQTWZC'
187 """
188 letter = unaccent(accented_letter)
189 if letter in string.ascii_letters:
190 if letter in string.ascii_uppercase:
191 alphabet_start = ord('A')
192 else:
193 alphabet_start = ord('a')
194 letter_number = ord(letter) - alphabet_start
195 if one_based: letter_number += 1
196 cipher_number = (letter_number * multiplier + adder) % 26
197 if one_based: cipher_number -= 1
198 return chr(cipher_number % 26 + alphabet_start)
199 else:
200 return letter
201
202 def affine_decipher_letter(letter, multiplier=1, adder=0, one_based=True):
203 """Encipher a letter, given a multiplier and adder
204
205 >>> ''.join([affine_decipher_letter(l, 3, 5, True) \
206 for l in 'HKNQTWZCFILORUXADGJMPSVYBE'])
207 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
208 >>> ''.join([affine_decipher_letter(l, 3, 5, False) \
209 for l in 'FILORUXADGJMPSVYBEHKNQTWZC'])
210 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
211 """
212 if letter in string.ascii_letters:
213 if letter in string.ascii_uppercase:
214 alphabet_start = ord('A')
215 else:
216 alphabet_start = ord('a')
217 cipher_number = ord(letter) - alphabet_start
218 if one_based: cipher_number += 1
219 plaintext_number = (
220 modular_division_table[multiplier]
221 [(cipher_number - adder) % 26] )
222 if one_based: plaintext_number -= 1
223 return chr(plaintext_number % 26 + alphabet_start)
224 else:
225 return letter
226
227 def affine_encipher(message, multiplier=1, adder=0, one_based=True):
228 """Encipher a message
229
230 >>> affine_encipher('hours passed during which jerico tried every ' \
231 'trick he could think of', 15, 22, True)
232 'lmyfu bkuusd dyfaxw claol psfaom jfasd snsfg jfaoe ls omytd jlaxe mh'
233 """
234 enciphered = [affine_encipher_letter(l, multiplier, adder, one_based)
235 for l in message]
236 return ''.join(enciphered)
237
238 def affine_decipher(message, multiplier=1, adder=0, one_based=True):
239 """Decipher a message
240
241 >>> affine_decipher('lmyfu bkuusd dyfaxw claol psfaom jfasd snsfg ' \
242 'jfaoe ls omytd jlaxe mh', 15, 22, True)
243 'hours passed during which jerico tried every trick he could think of'
244 """
245 enciphered = [affine_decipher_letter(l, multiplier, adder, one_based)
246 for l in message]
247 return ''.join(enciphered)
248
249
250 def keyword_cipher_alphabet_of(keyword, wrap_alphabet=0):
251 """Find the cipher alphabet given a keyword.
252 wrap_alphabet controls how the rest of the alphabet is added
253 after the keyword.
254 0 : from 'a'
255 1 : from the last letter in the sanitised keyword
256 2 : from the largest letter in the sanitised keyword
257
258 >>> keyword_cipher_alphabet_of('bayes')
259 'bayescdfghijklmnopqrtuvwxz'
260 >>> keyword_cipher_alphabet_of('bayes', 0)
261 'bayescdfghijklmnopqrtuvwxz'
262 >>> keyword_cipher_alphabet_of('bayes', 1)
263 'bayestuvwxzcdfghijklmnopqr'
264 >>> keyword_cipher_alphabet_of('bayes', 2)
265 'bayeszcdfghijklmnopqrtuvwx'
266 """
267 if wrap_alphabet == 0:
268 cipher_alphabet = ''.join(deduplicate(sanitise(keyword) +
269 string.ascii_lowercase))
270 else:
271 if wrap_alphabet == 1:
272 last_keyword_letter = deduplicate(sanitise(keyword))[-1]
273 else:
274 last_keyword_letter = sorted(sanitise(keyword))[-1]
275 last_keyword_position = string.ascii_lowercase.find(
276 last_keyword_letter) + 1
277 cipher_alphabet = ''.join(
278 deduplicate(sanitise(keyword) +
279 string.ascii_lowercase[last_keyword_position:] +
280 string.ascii_lowercase))
281 return cipher_alphabet
282
283
284 def keyword_encipher(message, keyword, wrap_alphabet=0):
285 """Enciphers a message with a keyword substitution cipher.
286 wrap_alphabet controls how the rest of the alphabet is added
287 after the keyword.
288 0 : from 'a'
289 1 : from the last letter in the sanitised keyword
290 2 : from the largest letter in the sanitised keyword
291
292 >>> keyword_encipher('test message', 'bayes')
293 'rsqr ksqqbds'
294 >>> keyword_encipher('test message', 'bayes', 0)
295 'rsqr ksqqbds'
296 >>> keyword_encipher('test message', 'bayes', 1)
297 'lskl dskkbus'
298 >>> keyword_encipher('test message', 'bayes', 2)
299 'qspq jsppbcs'
300 """
301 cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
302 cipher_translation = ''.maketrans(string.ascii_lowercase, cipher_alphabet)
303 return unaccent(message).lower().translate(cipher_translation)
304
305 def keyword_decipher(message, keyword, wrap_alphabet=0):
306 """Deciphers a message with a keyword substitution cipher.
307 wrap_alphabet controls how the rest of the alphabet is added
308 after the keyword.
309 0 : from 'a'
310 1 : from the last letter in the sanitised keyword
311 2 : from the largest letter in the sanitised keyword
312
313 >>> keyword_decipher('rsqr ksqqbds', 'bayes')
314 'test message'
315 >>> keyword_decipher('rsqr ksqqbds', 'bayes', 0)
316 'test message'
317 >>> keyword_decipher('lskl dskkbus', 'bayes', 1)
318 'test message'
319 >>> keyword_decipher('qspq jsppbcs', 'bayes', 2)
320 'test message'
321 """
322 cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
323 cipher_translation = ''.maketrans(cipher_alphabet, string.ascii_lowercase)
324 return message.lower().translate(cipher_translation)
325
326
327 def vigenere_encipher(message, keyword):
328 """Vigenere encipher
329
330 >>> vigenere_encipher('hello', 'abc')
331 'hfnlp'
332 """
333 shifts = [ord(l) - ord('a') for l in sanitise(keyword)]
334 pairs = zip(message, cycle(shifts))
335 return ''.join([caesar_encipher_letter(l, k) for l, k in pairs])
336
337 def vigenere_decipher(message, keyword):
338 """Vigenere decipher
339
340 >>> vigenere_decipher('hfnlp', 'abc')
341 'hello'
342 """
343 shifts = [ord(l) - ord('a') for l in sanitise(keyword)]
344 pairs = zip(message, cycle(shifts))
345 return ''.join([caesar_decipher_letter(l, k) for l, k in pairs])
346
347 beaufort_encipher=vigenere_decipher
348 beaufort_decipher=vigenere_encipher
349
350
351 def transpositions_of(keyword):
352 """Finds the transpostions given by a keyword. For instance, the keyword
353 'clever' rearranges to 'celrv', so the first column (0) stays first, the
354 second column (1) moves to third, the third column (2) moves to second,
355 and so on.
356
357 If passed a tuple, assume it's already a transposition and just return it.
358
359 >>> transpositions_of('clever')
360 (0, 2, 1, 4, 3)
361 >>> transpositions_of('fred')
362 (3, 2, 0, 1)
363 >>> transpositions_of((3, 2, 0, 1))
364 (3, 2, 0, 1)
365 """
366 if isinstance(keyword, tuple):
367 return keyword
368 else:
369 key = deduplicate(keyword)
370 transpositions = tuple(key.index(l) for l in sorted(key))
371 return transpositions
372
373 def pad(message_len, group_len, fillvalue):
374 padding_length = group_len - message_len % group_len
375 if padding_length == group_len: padding_length = 0
376 padding = ''
377 for i in range(padding_length):
378 if callable(fillvalue):
379 padding += fillvalue()
380 else:
381 padding += fillvalue
382 return padding
383
384 def column_transposition_encipher(message, keyword, fillvalue=' ',
385 fillcolumnwise=False,
386 emptycolumnwise=False):
387 """Enciphers using the column transposition cipher.
388 Message is padded to allow all rows to be the same length.
389
390 >>> column_transposition_encipher('hellothere', 'abcdef', fillcolumnwise=True)
391 'hlohr eltee '
392 >>> column_transposition_encipher('hellothere', 'abcdef', fillcolumnwise=True, emptycolumnwise=True)
393 'hellothere '
394 >>> column_transposition_encipher('hellothere', 'abcdef')
395 'hellothere '
396 >>> column_transposition_encipher('hellothere', 'abcde')
397 'hellothere'
398 >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=True)
399 'hellothere'
400 >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=False)
401 'hlohreltee'
402 >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=True)
403 'htehlelroe'
404 >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=False)
405 'hellothere'
406 >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=True, emptycolumnwise=True)
407 'heotllrehe'
408 >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=True, emptycolumnwise=False)
409 'holrhetlee'
410 >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=False, emptycolumnwise=True)
411 'htleehoelr'
412 >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=False, emptycolumnwise=False)
413 'hleolteher'
414 >>> column_transposition_encipher('hellothere', 'cleverly')
415 'hleolthre e '
416 >>> column_transposition_encipher('hellothere', 'cleverly', fillvalue='!')
417 'hleolthre!e!'
418 >>> column_transposition_encipher('hellothere', 'cleverly', fillvalue=lambda: '*')
419 'hleolthre*e*'
420 """
421 transpositions = transpositions_of(keyword)
422 message += pad(len(message), len(transpositions), fillvalue)
423 if fillcolumnwise:
424 rows = every_nth(message, len(message) // len(transpositions))
425 else:
426 rows = chunks(message, len(transpositions))
427 transposed = [transpose(r, transpositions) for r in rows]
428 if emptycolumnwise:
429 return combine_every_nth(transposed)
430 else:
431 return ''.join(chain(*transposed))
432
433 def column_transposition_decipher(message, keyword, fillvalue=' ',
434 fillcolumnwise=False,
435 emptycolumnwise=False):
436 """Deciphers using the column transposition cipher.
437 Message is padded to allow all rows to be the same length.
438
439 >>> column_transposition_decipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=True)
440 'hellothere'
441 >>> column_transposition_decipher('hlohreltee', 'abcde', fillcolumnwise=True, emptycolumnwise=False)
442 'hellothere'
443 >>> column_transposition_decipher('htehlelroe', 'abcde', fillcolumnwise=False, emptycolumnwise=True)
444 'hellothere'
445 >>> column_transposition_decipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=False)
446 'hellothere'
447 >>> column_transposition_decipher('heotllrehe', 'clever', fillcolumnwise=True, emptycolumnwise=True)
448 'hellothere'
449 >>> column_transposition_decipher('holrhetlee', 'clever', fillcolumnwise=True, emptycolumnwise=False)
450 'hellothere'
451 >>> column_transposition_decipher('htleehoelr', 'clever', fillcolumnwise=False, emptycolumnwise=True)
452 'hellothere'
453 >>> column_transposition_decipher('hleolteher', 'clever', fillcolumnwise=False, emptycolumnwise=False)
454 'hellothere'
455 """
456 transpositions = transpositions_of(keyword)
457 message += pad(len(message), len(transpositions), '*')
458 if emptycolumnwise:
459 rows = every_nth(message, len(message) // len(transpositions))
460 else:
461 rows = chunks(message, len(transpositions))
462 untransposed = [untranspose(r, transpositions) for r in rows]
463 if fillcolumnwise:
464 return combine_every_nth(untransposed)
465 else:
466 return ''.join(chain(*untransposed))
467
468 def scytale_encipher(message, rows, fillvalue=' '):
469 """Enciphers using the scytale transposition cipher.
470 Message is padded with spaces to allow all rows to be the same length.
471
472 >>> scytale_encipher('thequickbrownfox', 3)
473 'tcnhkfeboqrxuo iw '
474 >>> scytale_encipher('thequickbrownfox', 4)
475 'tubnhirfecooqkwx'
476 >>> scytale_encipher('thequickbrownfox', 5)
477 'tubnhirfecooqkwx'
478 >>> scytale_encipher('thequickbrownfox', 6)
479 'tqcrnxhukof eibwo '
480 >>> scytale_encipher('thequickbrownfox', 7)
481 'tqcrnxhukof eibwo '
482 """
483 transpositions = [i for i in range(math.ceil(len(message) / rows))]
484 return column_transposition_encipher(message, transpositions,
485 fillcolumnwise=False, emptycolumnwise=True)
486
487 def scytale_decipher(message, rows):
488 """Deciphers using the scytale transposition cipher.
489 Assumes the message is padded so that all rows are the same length.
490
491 >>> scytale_decipher('tcnhkfeboqrxuo iw ', 3)
492 'thequickbrownfox '
493 >>> scytale_decipher('tubnhirfecooqkwx', 4)
494 'thequickbrownfox'
495 >>> scytale_decipher('tubnhirfecooqkwx', 5)
496 'thequickbrownfox'
497 >>> scytale_decipher('tqcrnxhukof eibwo ', 6)
498 'thequickbrownfox '
499 >>> scytale_decipher('tqcrnxhukof eibwo ', 7)
500 'thequickbrownfox '
501 """
502 transpositions = [i for i in range(math.ceil(len(message) / rows))]
503 return column_transposition_decipher(message, transpositions,
504 fillcolumnwise=False, emptycolumnwise=True)
505
506
507 if __name__ == "__main__":
508 import doctest
509 doctest.testmod()