cipher.py

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