cipher.py

   1 import string
   2 import collections
   3 import math
   4 from enum import Enum
   5 from itertools import zip_longest, cycle, chain, count
   6 import numpy as np
   7 from numpy import matrix
   8 from numpy import linalg
   9 from language_models import *
  10
  11
  12 ## Utility functions
  13 cat = ''.join
  14 wcat = ' '.join
  15
  16
  17 modular_division_table = [[0]*26 for _ in range(26)]
  18 for a in range(26):
  19     for b in range(26):
  20         c = (a * b) % 26
  21         modular_division_table[b][c] = a
  22
  23
  24 def every_nth(text, n, fillvalue=''):
  25     """Returns n strings, each of which consists of every nth character,
  26     starting with the 0th, 1st, 2nd, ... (n-1)th character
  27
  28     >>> every_nth(string.ascii_lowercase, 5)
  29     ['afkpuz', 'bglqv', 'chmrw', 'dinsx', 'ejoty']
  30     >>> every_nth(string.ascii_lowercase, 1)
  31     ['abcdefghijklmnopqrstuvwxyz']
  32     >>> every_nth(string.ascii_lowercase, 26) # doctest: +NORMALIZE_WHITESPACE
  33     ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
  34      'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
  35     >>> every_nth(string.ascii_lowercase, 5, fillvalue='!')
  36     ['afkpuz', 'bglqv!', 'chmrw!', 'dinsx!', 'ejoty!']
  37     """
  38     split_text = chunks(text, n, fillvalue)
  39     return [cat(l) for l in zip_longest(*split_text, fillvalue=fillvalue)]
  40
  41 def combine_every_nth(split_text):
  42     """Reforms a text split into every_nth strings
  43
  44     >>> combine_every_nth(every_nth(string.ascii_lowercase, 5))
  45     'abcdefghijklmnopqrstuvwxyz'
  46     >>> combine_every_nth(every_nth(string.ascii_lowercase, 1))
  47     'abcdefghijklmnopqrstuvwxyz'
  48     >>> combine_every_nth(every_nth(string.ascii_lowercase, 26))
  49     'abcdefghijklmnopqrstuvwxyz'
  50     """
  51     return cat([cat(l)
  52                     for l in zip_longest(*split_text, fillvalue='')])
  53
  54 def chunks(text, n, fillvalue=None):
  55     """Split a text into chunks of n characters
  56
  57     >>> chunks('abcdefghi', 3)
  58     ['abc', 'def', 'ghi']
  59     >>> chunks('abcdefghi', 4)
  60     ['abcd', 'efgh', 'i']
  61     >>> chunks('abcdefghi', 4, fillvalue='!')
  62     ['abcd', 'efgh', 'i!!!']
  63     """
  64     if fillvalue:
  65         padding = fillvalue[0] * (n - len(text) % n)
  66     else:
  67         padding = ''
  68     return [(text+padding)[i:i+n] for i in range(0, len(text), n)]
  69
  70 def transpose(items, transposition):
  71     """Moves items around according to the given transposition
  72
  73     >>> transpose(['a', 'b', 'c', 'd'], (0,1,2,3))
  74     ['a', 'b', 'c', 'd']
  75     >>> transpose(['a', 'b', 'c', 'd'], (3,1,2,0))
  76     ['d', 'b', 'c', 'a']
  77     >>> transpose([10,11,12,13,14,15], (3,2,4,1,5,0))
  78     [13, 12, 14, 11, 15, 10]
  79     """
  80     transposed = [''] * len(transposition)
  81     for p, t in enumerate(transposition):
  82        transposed[p] = items[t]
  83     return transposed
  84
  85 def untranspose(items, transposition):
  86     """Undoes a transpose
  87
  88     >>> untranspose(['a', 'b', 'c', 'd'], [0,1,2,3])
  89     ['a', 'b', 'c', 'd']
  90     >>> untranspose(['d', 'b', 'c', 'a'], [3,1,2,0])
  91     ['a', 'b', 'c', 'd']
  92     >>> untranspose([13, 12, 14, 11, 15, 10], [3,2,4,1,5,0])
  93     [10, 11, 12, 13, 14, 15]
  94     """
  95     transposed = [''] * len(transposition)
  96     for p, t in enumerate(transposition):
  97        transposed[t] = items[p]
  98     return transposed
  99
 100 def deduplicate(text):
 101     return list(collections.OrderedDict.fromkeys(text))
 102
 103
 104 def caesar_encipher_letter(accented_letter, shift):
 105     """Encipher a letter, given a shift amount
 106
 107     >>> caesar_encipher_letter('a', 1)
 108     'b'
 109     >>> caesar_encipher_letter('a', 2)
 110     'c'
 111     >>> caesar_encipher_letter('b', 2)
 112     'd'
 113     >>> caesar_encipher_letter('x', 2)
 114     'z'
 115     >>> caesar_encipher_letter('y', 2)
 116     'a'
 117     >>> caesar_encipher_letter('z', 2)
 118     'b'
 119     >>> caesar_encipher_letter('z', -1)
 120     'y'
 121     >>> caesar_encipher_letter('a', -1)
 122     'z'
 123     >>> caesar_encipher_letter('A', 1)
 124     'B'
 125     >>> caesar_encipher_letter('é', 1)
 126     'f'
 127     """
 128     letter = unaccent(accented_letter)
 129     if letter in string.ascii_letters:
 130         if letter in string.ascii_uppercase:
 131             alphabet_start = ord('A')
 132         else:
 133             alphabet_start = ord('a')
 134         return chr(((ord(letter) - alphabet_start + shift) % 26) +
 135                    alphabet_start)
 136     else:
 137         return letter
 138
 139 def caesar_decipher_letter(letter, shift):
 140     """Decipher a letter, given a shift amount
 141
 142     >>> caesar_decipher_letter('b', 1)
 143     'a'
 144     >>> caesar_decipher_letter('b', 2)
 145     'z'
 146     """
 147     return caesar_encipher_letter(letter, -shift)
 148
 149 def caesar_encipher(message, shift):
 150     """Encipher a message with the Caesar cipher of given shift
 151
 152     >>> caesar_encipher('abc', 1)
 153     'bcd'
 154     >>> caesar_encipher('abc', 2)
 155     'cde'
 156     >>> caesar_encipher('abcxyz', 2)
 157     'cdezab'
 158     >>> caesar_encipher('ab cx yz', 2)
 159     'cd ez ab'
 160     >>> caesar_encipher('Héllo World!', 2)
 161     'Jgnnq Yqtnf!'
 162     """
 163     enciphered = [caesar_encipher_letter(l, shift) for l in message]
 164     return cat(enciphered)
 165
 166 def caesar_decipher(message, shift):
 167     """Decipher a message with the Caesar cipher of given shift
 168
 169     >>> caesar_decipher('bcd', 1)
 170     'abc'
 171     >>> caesar_decipher('cde', 2)
 172     'abc'
 173     >>> caesar_decipher('cd ez ab', 2)
 174     'ab cx yz'
 175     >>> caesar_decipher('Jgnnq Yqtnf!', 2)
 176     'Hello World!'
 177     """
 178     return caesar_encipher(message, -shift)
 179
 180 def affine_encipher_letter(accented_letter, multiplier=1, adder=0, one_based=True):
 181     """Encipher a letter, given a multiplier and adder
 182
 183     >>> cat([affine_encipher_letter(l, 3, 5, True) \
 184             for l in string.ascii_uppercase])
 185     'HKNQTWZCFILORUXADGJMPSVYBE'
 186     >>> cat([affine_encipher_letter(l, 3, 5, False) \
 187             for l in string.ascii_uppercase])
 188     'FILORUXADGJMPSVYBEHKNQTWZC'
 189     """
 190     letter = unaccent(accented_letter)
 191     if letter in string.ascii_letters:
 192         if letter in string.ascii_uppercase:
 193             alphabet_start = ord('A')
 194         else:
 195             alphabet_start = ord('a')
 196         letter_number = ord(letter) - alphabet_start
 197         if one_based: letter_number += 1
 198         cipher_number = (letter_number * multiplier + adder) % 26
 199         if one_based: cipher_number -= 1
 200         return chr(cipher_number % 26 + alphabet_start)
 201     else:
 202         return letter
 203
 204 def affine_decipher_letter(letter, multiplier=1, adder=0, one_based=True):
 205     """Encipher a letter, given a multiplier and adder
 206
 207     >>> cat([affine_decipher_letter(l, 3, 5, True) \
 208             for l in 'HKNQTWZCFILORUXADGJMPSVYBE'])
 209     'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
 210     >>> cat([affine_decipher_letter(l, 3, 5, False) \
 211             for l in 'FILORUXADGJMPSVYBEHKNQTWZC'])
 212     'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
 213     """
 214     if letter in string.ascii_letters:
 215         if letter in string.ascii_uppercase:
 216             alphabet_start = ord('A')
 217         else:
 218             alphabet_start = ord('a')
 219         cipher_number = ord(letter) - alphabet_start
 220         if one_based: cipher_number += 1
 221         plaintext_number = (
 222             modular_division_table[multiplier]
 223                                   [(cipher_number - adder) % 26])
 224         if one_based: plaintext_number -= 1
 225         return chr(plaintext_number % 26 + alphabet_start)
 226     else:
 227         return letter
 228
 229 def affine_encipher(message, multiplier=1, adder=0, one_based=True):
 230     """Encipher a message
 231
 232     >>> affine_encipher('hours passed during which jerico tried every ' \
 233            'trick he could think of', 15, 22, True)
 234     'lmyfu bkuusd dyfaxw claol psfaom jfasd snsfg jfaoe ls omytd jlaxe mh'
 235     """
 236     enciphered = [affine_encipher_letter(l, multiplier, adder, one_based)
 237                   for l in message]
 238     return cat(enciphered)
 239
 240 def affine_decipher(message, multiplier=1, adder=0, one_based=True):
 241     """Decipher a message
 242
 243     >>> affine_decipher('lmyfu bkuusd dyfaxw claol psfaom jfasd snsfg ' \
 244            'jfaoe ls omytd jlaxe mh', 15, 22, True)
 245     'hours passed during which jerico tried every trick he could think of'
 246     """
 247     enciphered = [affine_decipher_letter(l, multiplier, adder, one_based)
 248                   for l in message]
 249     return cat(enciphered)
 250
 251
 252 class KeywordWrapAlphabet(Enum):
 253     from_a = 1
 254     from_last = 2
 255     from_largest = 3
 256
 257
 258 def keyword_cipher_alphabet_of(keyword, wrap_alphabet=KeywordWrapAlphabet.from_a):
 259     """Find the cipher alphabet given a keyword.
 260     wrap_alphabet controls how the rest of the alphabet is added
 261     after the keyword.
 262
 263     >>> keyword_cipher_alphabet_of('bayes')
 264     'bayescdfghijklmnopqrtuvwxz'
 265     >>> keyword_cipher_alphabet_of('bayes', KeywordWrapAlphabet.from_a)
 266     'bayescdfghijklmnopqrtuvwxz'
 267     >>> keyword_cipher_alphabet_of('bayes', KeywordWrapAlphabet.from_last)
 268     'bayestuvwxzcdfghijklmnopqr'
 269     >>> keyword_cipher_alphabet_of('bayes', KeywordWrapAlphabet.from_largest)
 270     'bayeszcdfghijklmnopqrtuvwx'
 271     """
 272     if wrap_alphabet == KeywordWrapAlphabet.from_a:
 273         cipher_alphabet = cat(deduplicate(sanitise(keyword) +
 274                                               string.ascii_lowercase))
 275     else:
 276         if wrap_alphabet == KeywordWrapAlphabet.from_last:
 277             last_keyword_letter = deduplicate(sanitise(keyword))[-1]
 278         else:
 279             last_keyword_letter = sorted(sanitise(keyword))[-1]
 280         last_keyword_position = string.ascii_lowercase.find(
 281             last_keyword_letter) + 1
 282         cipher_alphabet = cat(
 283             deduplicate(sanitise(keyword) +
 284                         string.ascii_lowercase[last_keyword_position:] +
 285                         string.ascii_lowercase))
 286     return cipher_alphabet
 287
 288
 289 def keyword_encipher(message, keyword, wrap_alphabet=KeywordWrapAlphabet.from_a):
 290     """Enciphers a message with a keyword substitution cipher.
 291     wrap_alphabet controls how the rest of the alphabet is added
 292     after the keyword.
 293     0 : from 'a'
 294     1 : from the last letter in the sanitised keyword
 295     2 : from the largest letter in the sanitised keyword
 296
 297     >>> keyword_encipher('test message', 'bayes')
 298     'rsqr ksqqbds'
 299     >>> keyword_encipher('test message', 'bayes', KeywordWrapAlphabet.from_a)
 300     'rsqr ksqqbds'
 301     >>> keyword_encipher('test message', 'bayes', KeywordWrapAlphabet.from_last)
 302     'lskl dskkbus'
 303     >>> keyword_encipher('test message', 'bayes', KeywordWrapAlphabet.from_largest)
 304     'qspq jsppbcs'
 305     """
 306     cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
 307     cipher_translation = ''.maketrans(string.ascii_lowercase, cipher_alphabet)
 308     return unaccent(message).lower().translate(cipher_translation)
 309
 310 def keyword_decipher(message, keyword, wrap_alphabet=KeywordWrapAlphabet.from_a):
 311     """Deciphers a message with a keyword substitution cipher.
 312     wrap_alphabet controls how the rest of the alphabet is added
 313     after the keyword.
 314     0 : from 'a'
 315     1 : from the last letter in the sanitised keyword
 316     2 : from the largest letter in the sanitised keyword
 317
 318     >>> keyword_decipher('rsqr ksqqbds', 'bayes')
 319     'test message'
 320     >>> keyword_decipher('rsqr ksqqbds', 'bayes', KeywordWrapAlphabet.from_a)
 321     'test message'
 322     >>> keyword_decipher('lskl dskkbus', 'bayes', KeywordWrapAlphabet.from_last)
 323     'test message'
 324     >>> keyword_decipher('qspq jsppbcs', 'bayes', KeywordWrapAlphabet.from_largest)
 325     'test message'
 326     """
 327     cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
 328     cipher_translation = ''.maketrans(cipher_alphabet, string.ascii_lowercase)
 329     return message.lower().translate(cipher_translation)
 330
 331
 332 def vigenere_encipher(message, keyword):
 333     """Vigenere encipher
 334
 335     >>> vigenere_encipher('hello', 'abc')
 336     'hfnlp'
 337     """
 338     shifts = [ord(l) - ord('a') for l in sanitise(keyword)]
 339     pairs = zip(message, cycle(shifts))
 340     return cat([caesar_encipher_letter(l, k) for l, k in pairs])
 341
 342 def vigenere_decipher(message, keyword):
 343     """Vigenere decipher
 344
 345     >>> vigenere_decipher('hfnlp', 'abc')
 346     'hello'
 347     """
 348     shifts = [ord(l) - ord('a') for l in sanitise(keyword)]
 349     pairs = zip(message, cycle(shifts))
 350     return cat([caesar_decipher_letter(l, k) for l, k in pairs])
 351
 352 beaufort_encipher=vigenere_decipher
 353 beaufort_decipher=vigenere_encipher
 354
 355
 356 def transpositions_of(keyword):
 357     """Finds the transpostions given by a keyword. For instance, the keyword
 358     'clever' rearranges to 'celrv', so the first column (0) stays first, the
 359     second column (1) moves to third, the third column (2) moves to second,
 360     and so on.
 361
 362     If passed a tuple, assume it's already a transposition and just return it.
 363
 364     >>> transpositions_of('clever')
 365     (0, 2, 1, 4, 3)
 366     >>> transpositions_of('fred')
 367     (3, 2, 0, 1)
 368     >>> transpositions_of((3, 2, 0, 1))
 369     (3, 2, 0, 1)
 370     """
 371     if isinstance(keyword, tuple):
 372         return keyword
 373     else:
 374         key = deduplicate(keyword)
 375         transpositions = tuple(key.index(l) for l in sorted(key))
 376         return transpositions
 377
 378 def pad(message_len, group_len, fillvalue):
 379     padding_length = group_len - message_len % group_len
 380     if padding_length == group_len: padding_length = 0
 381     padding = ''
 382     for i in range(padding_length):
 383         if callable(fillvalue):
 384             padding += fillvalue()
 385         else:
 386             padding += fillvalue
 387     return padding
 388
 389 def column_transposition_encipher(message, keyword, fillvalue=' ',
 390       fillcolumnwise=False,
 391       emptycolumnwise=False):
 392     """Enciphers using the column transposition cipher.
 393     Message is padded to allow all rows to be the same length.
 394
 395     >>> column_transposition_encipher('hellothere', 'abcdef', fillcolumnwise=True)
 396     'hlohr eltee '
 397     >>> column_transposition_encipher('hellothere', 'abcdef', fillcolumnwise=True, emptycolumnwise=True)
 398     'hellothere  '
 399     >>> column_transposition_encipher('hellothere', 'abcdef')
 400     'hellothere  '
 401     >>> column_transposition_encipher('hellothere', 'abcde')
 402     'hellothere'
 403     >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=True)
 404     'hellothere'
 405     >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=False)
 406     'hlohreltee'
 407     >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=True)
 408     'htehlelroe'
 409     >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=False)
 410     'hellothere'
 411     >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=True, emptycolumnwise=True)
 412     'heotllrehe'
 413     >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=True, emptycolumnwise=False)
 414     'holrhetlee'
 415     >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=False, emptycolumnwise=True)
 416     'htleehoelr'
 417     >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=False, emptycolumnwise=False)
 418     'hleolteher'
 419     >>> column_transposition_encipher('hellothere', 'cleverly')
 420     'hleolthre e '
 421     >>> column_transposition_encipher('hellothere', 'cleverly', fillvalue='!')
 422     'hleolthre!e!'
 423     >>> column_transposition_encipher('hellothere', 'cleverly', fillvalue=lambda: '*')
 424     'hleolthre*e*'
 425     """
 426     transpositions = transpositions_of(keyword)
 427     message += pad(len(message), len(transpositions), fillvalue)
 428     if fillcolumnwise:
 429         rows = every_nth(message, len(message) // len(transpositions))
 430     else:
 431         rows = chunks(message, len(transpositions))
 432     transposed = [transpose(r, transpositions) for r in rows]
 433     if emptycolumnwise:
 434         return combine_every_nth(transposed)
 435     else:
 436         return cat(chain(*transposed))
 437
 438 def column_transposition_decipher(message, keyword, fillvalue=' ',
 439       fillcolumnwise=False,
 440       emptycolumnwise=False):
 441     """Deciphers using the column transposition cipher.
 442     Message is padded to allow all rows to be the same length.
 443
 444     >>> column_transposition_decipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=True)
 445     'hellothere'
 446     >>> column_transposition_decipher('hlohreltee', 'abcde', fillcolumnwise=True, emptycolumnwise=False)
 447     'hellothere'
 448     >>> column_transposition_decipher('htehlelroe', 'abcde', fillcolumnwise=False, emptycolumnwise=True)
 449     'hellothere'
 450     >>> column_transposition_decipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=False)
 451     'hellothere'
 452     >>> column_transposition_decipher('heotllrehe', 'clever', fillcolumnwise=True, emptycolumnwise=True)
 453     'hellothere'
 454     >>> column_transposition_decipher('holrhetlee', 'clever', fillcolumnwise=True, emptycolumnwise=False)
 455     'hellothere'
 456     >>> column_transposition_decipher('htleehoelr', 'clever', fillcolumnwise=False, emptycolumnwise=True)
 457     'hellothere'
 458     >>> column_transposition_decipher('hleolteher', 'clever', fillcolumnwise=False, emptycolumnwise=False)
 459     'hellothere'
 460     """
 461     transpositions = transpositions_of(keyword)
 462     message += pad(len(message), len(transpositions), fillvalue)
 463     if emptycolumnwise:
 464         rows = every_nth(message, len(message) // len(transpositions))
 465     else:
 466         rows = chunks(message, len(transpositions))
 467     untransposed = [untranspose(r, transpositions) for r in rows]
 468     if fillcolumnwise:
 469         return combine_every_nth(untransposed)
 470     else:
 471         return cat(chain(*untransposed))
 472
 473 def scytale_encipher(message, rows, fillvalue=' '):
 474     """Enciphers using the scytale transposition cipher.
 475     Message is padded with spaces to allow all rows to be the same length.
 476
 477     >>> scytale_encipher('thequickbrownfox', 3)
 478     'tcnhkfeboqrxuo iw '
 479     >>> scytale_encipher('thequickbrownfox', 4)
 480     'tubnhirfecooqkwx'
 481     >>> scytale_encipher('thequickbrownfox', 5)
 482     'tubn hirf ecoo qkwx '
 483     >>> scytale_encipher('thequickbrownfox', 6)
 484     'tqcrnxhukof eibwo '
 485     >>> scytale_encipher('thequickbrownfox', 7)
 486     'tqcrnx hukof  eibwo  '
 487     """
 488     # transpositions = [i for i in range(math.ceil(len(message) / rows))]
 489     # return column_transposition_encipher(message, transpositions,
 490     #     fillvalue=fillvalue, fillcolumnwise=False, emptycolumnwise=True)
 491     transpositions = [i for i in range(rows)]
 492     return column_transposition_encipher(message, transpositions,
 493         fillvalue=fillvalue, fillcolumnwise=True, emptycolumnwise=False)
 494
 495 def scytale_decipher(message, rows):
 496     """Deciphers using the scytale transposition cipher.
 497     Assumes the message is padded so that all rows are the same length.
 498
 499     >>> scytale_decipher('tcnhkfeboqrxuo iw ', 3)
 500     'thequickbrownfox  '
 501     >>> scytale_decipher('tubnhirfecooqkwx', 4)
 502     'thequickbrownfox'
 503     >>> scytale_decipher('tubn hirf ecoo qkwx ', 5)
 504     'thequickbrownfox    '
 505     >>> scytale_decipher('tqcrnxhukof eibwo ', 6)
 506     'thequickbrownfox  '
 507     >>> scytale_decipher('tqcrnx hukof  eibwo  ', 7)
 508     'thequickbrownfox     '
 509     """
 510     # transpositions = [i for i in range(math.ceil(len(message) / rows))]
 511     # return column_transposition_decipher(message, transpositions,
 512     #     fillcolumnwise=False, emptycolumnwise=True)
 513     transpositions = [i for i in range(rows)]
 514     return column_transposition_decipher(message, transpositions,
 515         fillcolumnwise=True, emptycolumnwise=False)
 516
 517
 518 def railfence_encipher(message, height, fillvalue=''):
 519     """Railfence cipher.
 520     Works by splitting the text into sections, then reading across them to
 521     generate the rows in the cipher. The rows are then combined to form the
 522     ciphertext.
 523
 524     Example: the plaintext "hellotherefriends", with a height of four, written
 525     out in the railfence as
 526        h h i
 527        etere*
 528        lorfns
 529        l e d
 530     (with the * showing the one character to finish the last section).
 531     Each 'section' is two columns, but unfolded. In the example, the first
 532     section is 'hellot'.
 533
 534     >>> railfence_encipher('hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers', 2, fillvalue='!')
 535     'hlohraateerishsslnpeefetotsigaleccpeselteevsmhatetiiaogicotxfretnrifneihr!'
 536     >>> railfence_encipher('hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers', 3, fillvalue='!')
 537     'horaersslpeeosglcpselteevsmhatetiiaogicotxfretnrifneihr!!lhateihsnefttiaece!'
 538     >>> railfence_encipher('hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers', 5, fillvalue='!')
 539     'hresleogcseeemhetaocofrnrner!!lhateihsnefttiaece!!ltvsatiigitxetifih!!oarspeslp!'
 540     >>> railfence_encipher('hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers', 10, fillvalue='!')
 541     'hepisehagitnr!!lernesge!!lmtocerh!!otiletap!!tseaorii!!hassfolc!!evtitffe!!rahsetec!!eixn!'
 542     >>> railfence_encipher('hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers', 3)
 543     'horaersslpeeosglcpselteevsmhatetiiaogicotxfretnrifneihrlhateihsnefttiaece'
 544     >>> railfence_encipher('hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers', 5)
 545     'hresleogcseeemhetaocofrnrnerlhateihsnefttiaeceltvsatiigitxetifihoarspeslp'
 546     >>> railfence_encipher('hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers', 7)
 547     'haspolsevsetgifrifrlatihnettaeelemtiocxernhorersleesgcptehaiaottneihesfic'
 548     """
 549     sections = chunks(message, (height - 1) * 2, fillvalue=fillvalue)
 550     n_sections = len(sections)
 551     # Add the top row
 552     rows = [cat([s[0] for s in sections])]
 553     # process the middle rows of the grid
 554     for r in range(1, height-1):
 555         rows += [cat([s[r:r+1] + s[height*2-r-2:height*2-r-1] for s in sections])]
 556     # process the bottom row
 557     rows += [cat([s[height - 1:height] for s in sections])]
 558     # rows += [wcat([s[height - 1] for s in sections])]
 559     return cat(rows)
 560
 561 def railfence_decipher(message, height, fillvalue=''):
 562     """Railfence decipher.
 563     Works by reconstructing the grid used to generate the ciphertext, then
 564     unfolding the sections so the text can be concatenated together.
 565
 566     Example: given the ciphertext 'hhieterelorfnsled' and a height of 4, first
 567     work out that the second row has a character missing, find the rows of the
 568     grid, then split the section into its two columns.
 569
 570     'hhieterelorfnsled' is split into
 571         h h i
 572         etere
 573         lorfns
 574         l e d
 575     (spaces added for clarity), which is stored in 'rows'. This is then split
 576     into 'down_rows' and 'up_rows':
 577
 578     down_rows:
 579        hhi
 580        eee
 581        lrn
 582        led
 583
 584     up_rows:
 585        tr
 586        ofs
 587
 588     These are then zipped together (after the up_rows are reversed) to recover
 589     the plaintext.
 590
 591     Most of the procedure is about finding the correct lengths for each row then
 592     splitting the ciphertext into those rows.
 593
 594     >>> railfence_decipher('hlohraateerishsslnpeefetotsigaleccpeselteevsmhatetiiaogicotxfretnrifneihr!', 2).strip('!')
 595     'hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers'
 596     >>> railfence_decipher('horaersslpeeosglcpselteevsmhatetiiaogicotxfretnrifneihr!!lhateihsnefttiaece!', 3).strip('!')
 597     'hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers'
 598     >>> railfence_decipher('hresleogcseeemhetaocofrnrner!!lhateihsnefttiaece!!ltvsatiigitxetifih!!oarspeslp!', 5).strip('!')
 599     'hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers'
 600     >>> railfence_decipher('hepisehagitnr!!lernesge!!lmtocerh!!otiletap!!tseaorii!!hassfolc!!evtitffe!!rahsetec!!eixn!', 10).strip('!')
 601     'hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers'
 602     >>> railfence_decipher('horaersslpeeosglcpselteevsmhatetiiaogicotxfretnrifneihrlhateihsnefttiaece', 3)
 603     'hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers'
 604     >>> railfence_decipher('hresleogcseeemhetaocofrnrnerlhateihsnefttiaeceltvsatiigitxetifihoarspeslp', 5)
 605     'hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers'
 606     >>> railfence_decipher('haspolsevsetgifrifrlatihnettaeelemtiocxernhorersleesgcptehaiaottneihesfic', 7)
 607     'hellothereavastmeheartiesthisisalongpieceoftextfortestingrailfenceciphers'
 608     """
 609     # find the number and size of the sections, including how many characters
 610     #   are missing for a full grid
 611     n_sections = math.ceil(len(message) / ((height - 1) * 2))
 612     padding_to_add = n_sections * (height - 1) * 2 - len(message)
 613     # row_lengths are for the both up rows and down rows
 614     row_lengths = [n_sections] * (height - 1) * 2
 615     for i in range((height - 1) * 2 - 1, (height - 1) * 2 - (padding_to_add + 1), -1):
 616         row_lengths[i] -= 1
 617     # folded_rows are the combined row lengths in the middle of the railfence
 618     folded_row_lengths = [row_lengths[0]]
 619     for i in range(1, height-1):
 620         folded_row_lengths += [row_lengths[i] + row_lengths[-i]]
 621     folded_row_lengths += [row_lengths[height - 1]]
 622     # find the rows that form the railfence grid
 623     rows = []
 624     row_start = 0
 625     for i in folded_row_lengths:
 626         rows += [message[row_start:row_start + i]]
 627         row_start += i
 628     # split the rows into the 'down_rows' (those that form the first column of
 629     #   a section) and the 'up_rows' (those that ofrm the second column of a
 630     #   section).
 631     down_rows = [rows[0]]
 632     up_rows = []
 633     for i in range(1, height-1):
 634         down_rows += [cat([c for n, c in enumerate(rows[i]) if n % 2 == 0])]
 635         up_rows += [cat([c for n, c in enumerate(rows[i]) if n % 2 == 1])]
 636     down_rows += [rows[-1]]
 637     up_rows.reverse()
 638     return cat(c for r in zip_longest(*(down_rows + up_rows), fillvalue='') for c in r)
 639
 640 def make_cadenus_keycolumn(doubled_letters = 'vw', start='a', reverse=False):
 641     """Makes the key column for a Cadenus cipher (the column down between the
 642         rows of letters)
 643
 644     >>> make_cadenus_keycolumn()['a']
 645     0
 646     >>> make_cadenus_keycolumn()['b']
 647     1
 648     >>> make_cadenus_keycolumn()['c']
 649     2
 650     >>> make_cadenus_keycolumn()['v']
 651     21
 652     >>> make_cadenus_keycolumn()['w']
 653     21
 654     >>> make_cadenus_keycolumn()['z']
 655     24
 656     >>> make_cadenus_keycolumn(doubled_letters='ij', start='b', reverse=True)['a']
 657     1
 658     >>> make_cadenus_keycolumn(doubled_letters='ij', start='b', reverse=True)['b']
 659     0
 660     >>> make_cadenus_keycolumn(doubled_letters='ij', start='b', reverse=True)['c']
 661     24
 662     >>> make_cadenus_keycolumn(doubled_letters='ij', start='b', reverse=True)['i']
 663     18
 664     >>> make_cadenus_keycolumn(doubled_letters='ij', start='b', reverse=True)['j']
 665     18
 666     >>> make_cadenus_keycolumn(doubled_letters='ij', start='b', reverse=True)['v']
 667     6
 668     >>> make_cadenus_keycolumn(doubled_letters='ij', start='b', reverse=True)['z']
 669     2
 670     """
 671     index_to_remove = string.ascii_lowercase.find(doubled_letters[0])
 672     short_alphabet = string.ascii_lowercase[:index_to_remove] + string.ascii_lowercase[index_to_remove+1:]
 673     if reverse:
 674         short_alphabet = cat(reversed(short_alphabet))
 675     start_pos = short_alphabet.find(start)
 676     rotated_alphabet = short_alphabet[start_pos:] + short_alphabet[:start_pos]
 677     keycolumn = {l: i for i, l in enumerate(rotated_alphabet)}
 678     keycolumn[doubled_letters[0]] = keycolumn[doubled_letters[1]]
 679     return keycolumn
 680
 681 def cadenus_encipher(message, keyword, keycolumn, fillvalue='a'):
 682     """Encipher with the Cadenus cipher
 683
 684     >>> cadenus_encipher(sanitise('Whoever has made a voyage up the Hudson ' \
 685                                   'must remember the Kaatskill mountains. ' \
 686                                   'They are a dismembered branch of the great'), \
 687                 'wink', \
 688                 make_cadenus_keycolumn(doubled_letters='vw', start='a', reverse=True))
 689     'antodeleeeuhrsidrbhmhdrrhnimefmthgeaetakseomehetyaasuvoyegrastmmuuaeenabbtpchehtarorikswosmvaleatned'
 690     >>> cadenus_encipher(sanitise('a severe limitation on the usefulness of ' \
 691                                   'the cadenus is that every message must be ' \
 692                                   'a multiple of twenty-five letters long'), \
 693                 'easy', \
 694                 make_cadenus_keycolumn(doubled_letters='vw', start='a', reverse=True))
 695     'systretomtattlusoatleeesfiyheasdfnmschbhneuvsnpmtofarenuseieeieltarlmentieetogevesitfaisltngeeuvowul'
 696     """
 697     rows = chunks(message, len(message) // 25, fillvalue=fillvalue)
 698     columns = zip(*rows)
 699     rotated_columns = [col[start:] + col[:start] for start, col in zip([keycolumn[l] for l in keyword], columns)]
 700     rotated_rows = zip(*rotated_columns)
 701     transpositions = transpositions_of(keyword)
 702     transposed = [transpose(r, transpositions) for r in rotated_rows]
 703     return cat(chain(*transposed))
 704
 705 def cadenus_decipher(message, keyword, keycolumn, fillvalue='a'):
 706     """
 707     >>> cadenus_decipher('antodeleeeuhrsidrbhmhdrrhnimefmthgeaetakseomehetyaa' \
 708                          'suvoyegrastmmuuaeenabbtpchehtarorikswosmvaleatned', \
 709                  'wink', \
 710                  make_cadenus_keycolumn(reverse=True))
 711     'whoeverhasmadeavoyageupthehudsonmustrememberthekaatskillmountainstheyareadismemberedbranchofthegreat'
 712     >>> cadenus_decipher('systretomtattlusoatleeesfiyheasdfnmschbhneuvsnpmtof' \
 713                         'arenuseieeieltarlmentieetogevesitfaisltngeeuvowul', \
 714                  'easy', \
 715                  make_cadenus_keycolumn(reverse=True))
 716     'aseverelimitationontheusefulnessofthecadenusisthateverymessagemustbeamultipleoftwentyfiveletterslong'
 717     """
 718     rows = chunks(message, len(message) // 25, fillvalue=fillvalue)
 719     transpositions = transpositions_of(keyword)
 720     untransposed_rows = [untranspose(r, transpositions) for r in rows]
 721     columns = zip(*untransposed_rows)
 722     rotated_columns = [col[-start:] + col[:-start] for start, col in zip([keycolumn[l] for l in keyword], columns)]
 723     rotated_rows = zip(*rotated_columns)
 724     # return rotated_columns
 725     return cat(chain(*rotated_rows))
 726
 727
 728 def hill_encipher(matrix, message_letters, fillvalue='a'):
 729     """Hill cipher
 730
 731     >>> hill_encipher(np.matrix([[7,8], [11,11]]), 'hellothere')
 732     'drjiqzdrvx'
 733     >>> hill_encipher(np.matrix([[6, 24, 1], [13, 16, 10], [20, 17, 15]]), \
 734         'hello there')
 735     'tfjflpznvyac'
 736     """
 737     n = len(matrix)
 738     sanitised_message = sanitise(message_letters)
 739     if len(sanitised_message) % n != 0:
 740         padding = fillvalue[0] * (n - len(sanitised_message) % n)
 741     else:
 742         padding = ''
 743     message = [ord(c) - ord('a') for c in sanitised_message + padding]
 744     message_chunks = [message[i:i+n] for i in range(0, len(message), n)]
 745     # message_chunks = chunks(message, len(matrix), fillvalue=None)
 746     enciphered_chunks = [((matrix * np.matrix(c).T).T).tolist()[0]
 747             for c in message_chunks]
 748     return cat([chr(int(round(l)) % 26 + ord('a'))
 749             for l in sum(enciphered_chunks, [])])
 750
 751 def hill_decipher(matrix, message, fillvalue='a'):
 752     """Hill cipher
 753
 754     >>> hill_decipher(np.matrix([[7,8], [11,11]]), 'drjiqzdrvx')
 755     'hellothere'
 756     >>> hill_decipher(np.matrix([[6, 24, 1], [13, 16, 10], [20, 17, 15]]), \
 757         'tfjflpznvyac')
 758     'hellothereaa'
 759     """
 760     adjoint = linalg.det(matrix)*linalg.inv(matrix)
 761     inverse_determinant = modular_division_table[int(round(linalg.det(matrix))) % 26][1]
 762     inverse_matrix = (inverse_determinant * adjoint) % 26
 763     return hill_encipher(inverse_matrix, message, fillvalue)
 764
 765
 766 # Where each piece of text ends up in the AMSCO transpositon cipher.
 767 # 'index' shows where the slice appears in the plaintext, with the slice
 768 # from 'start' to 'end'
 769 AmscoSlice = collections.namedtuple('AmscoSlice', ['index', 'start', 'end'])
 770
 771 class AmscoFillStyle(Enum):
 772     continuous = 1
 773     same_each_row = 2
 774     reverse_each_row = 3
 775
 776 def amsco_transposition_positions(message, keyword,
 777       fillpattern=(1, 2),
 778       fillstyle=AmscoFillStyle.continuous,
 779       fillcolumnwise=False,
 780       emptycolumnwise=True):
 781     """Creates the grid for the AMSCO transposition cipher. Each element in the
 782     grid shows the index of that slice and the start and end positions of the
 783     plaintext that go to make it up.
 784
 785     >>> amsco_transposition_positions(string.ascii_lowercase, 'freddy', \
 786         fillpattern=(1, 2)) # doctest:  +NORMALIZE_WHITESPACE
 787     [[AmscoSlice(index=3, start=4, end=6),
 788      AmscoSlice(index=2, start=3, end=4),
 789      AmscoSlice(index=0, start=0, end=1),
 790      AmscoSlice(index=1, start=1, end=3),
 791      AmscoSlice(index=4, start=6, end=7)],
 792     [AmscoSlice(index=8, start=12, end=13),
 793      AmscoSlice(index=7, start=10, end=12),
 794      AmscoSlice(index=5, start=7, end=9),
 795      AmscoSlice(index=6, start=9, end=10),
 796      AmscoSlice(index=9, start=13, end=15)],
 797     [AmscoSlice(index=13, start=19, end=21),
 798      AmscoSlice(index=12, start=18, end=19),
 799      AmscoSlice(index=10, start=15, end=16),
 800      AmscoSlice(index=11, start=16, end=18),
 801      AmscoSlice(index=14, start=21, end=22)],
 802     [AmscoSlice(index=18, start=27, end=28),
 803      AmscoSlice(index=17, start=25, end=27),
 804      AmscoSlice(index=15, start=22, end=24),
 805      AmscoSlice(index=16, start=24, end=25),
 806      AmscoSlice(index=19, start=28, end=30)]]
 807     """
 808     transpositions = transpositions_of(keyword)
 809     fill_iterator = cycle(fillpattern)
 810     indices = count()
 811     message_length = len(message)
 812
 813     current_position = 0
 814     grid = []
 815     current_fillpattern = fillpattern
 816     while current_position < message_length:
 817         row = []
 818         if fillstyle == AmscoFillStyle.same_each_row:
 819             fill_iterator = cycle(fillpattern)
 820         if fillstyle == AmscoFillStyle.reverse_each_row:
 821             fill_iterator = cycle(current_fillpattern)
 822         for _ in range(len(transpositions)):
 823             index = next(indices)
 824             gap = next(fill_iterator)
 825             row += [AmscoSlice(index, current_position, current_position + gap)]
 826             current_position += gap
 827         grid += [row]
 828         if fillstyle == AmscoFillStyle.reverse_each_row:
 829             current_fillpattern = list(reversed(current_fillpattern))
 830     return [transpose(r, transpositions) for r in grid]
 831
 832 def amsco_transposition_encipher(message, keyword,
 833     fillpattern=(1,2), fillstyle=AmscoFillStyle.reverse_each_row):
 834     """AMSCO transposition encipher.
 835
 836     >>> amsco_transposition_encipher('hellothere', 'abc', fillpattern=(1, 2))
 837     'hoteelhler'
 838     >>> amsco_transposition_encipher('hellothere', 'abc', fillpattern=(2, 1))
 839     'hetelhelor'
 840     >>> amsco_transposition_encipher('hellothere', 'acb', fillpattern=(1, 2))
 841     'hotelerelh'
 842     >>> amsco_transposition_encipher('hellothere', 'acb', fillpattern=(2, 1))
 843     'hetelorlhe'
 844     >>> amsco_transposition_encipher('hereissometexttoencipher', 'encode')
 845     'etecstthhomoerereenisxip'
 846     >>> amsco_transposition_encipher('hereissometexttoencipher', 'cipher', fillpattern=(1, 2))
 847     'hetcsoeisterereipexthomn'
 848     >>> amsco_transposition_encipher('hereissometexttoencipher', 'cipher', fillpattern=(1, 2), fillstyle=AmscoFillStyle.continuous)
 849     'hecsoisttererteipexhomen'
 850     >>> amsco_transposition_encipher('hereissometexttoencipher', 'cipher', fillpattern=(2, 1))
 851     'heecisoosttrrtepeixhemen'
 852     >>> amsco_transposition_encipher('hereissometexttoencipher', 'cipher', fillpattern=(1, 3, 2))
 853     'hxtomephescieretoeisnter'
 854     >>> amsco_transposition_encipher('hereissometexttoencipher', 'cipher', fillpattern=(1, 3, 2), fillstyle=AmscoFillStyle.continuous)
 855     'hxomeiphscerettoisenteer'
 856     """
 857     grid = amsco_transposition_positions(message, keyword,
 858         fillpattern=fillpattern, fillstyle=fillstyle)
 859     ct_as_grid = [[message[s.start:s.end] for s in r] for r in grid]
 860     return combine_every_nth(ct_as_grid)
 861
 862
 863 def amsco_transposition_decipher(message, keyword,
 864     fillpattern=(1,2), fillstyle=AmscoFillStyle.reverse_each_row):
 865     """AMSCO transposition decipher
 866
 867     >>> amsco_transposition_decipher('hoteelhler', 'abc', fillpattern=(1, 2))
 868     'hellothere'
 869     >>> amsco_transposition_decipher('hetelhelor', 'abc', fillpattern=(2, 1))
 870     'hellothere'
 871     >>> amsco_transposition_decipher('hotelerelh', 'acb', fillpattern=(1, 2))
 872     'hellothere'
 873     >>> amsco_transposition_decipher('hetelorlhe', 'acb', fillpattern=(2, 1))
 874     'hellothere'
 875     >>> amsco_transposition_decipher('etecstthhomoerereenisxip', 'encode')
 876     'hereissometexttoencipher'
 877     >>> amsco_transposition_decipher('hetcsoeisterereipexthomn', 'cipher', fillpattern=(1, 2))
 878     'hereissometexttoencipher'
 879     >>> amsco_transposition_decipher('hecsoisttererteipexhomen', 'cipher', fillpattern=(1, 2), fillstyle=AmscoFillStyle.continuous)
 880     'hereissometexttoencipher'
 881     >>> amsco_transposition_decipher('heecisoosttrrtepeixhemen', 'cipher', fillpattern=(2, 1))
 882     'hereissometexttoencipher'
 883     >>> amsco_transposition_decipher('hxtomephescieretoeisnter', 'cipher', fillpattern=(1, 3, 2))
 884     'hereissometexttoencipher'
 885     >>> amsco_transposition_decipher('hxomeiphscerettoisenteer', 'cipher', fillpattern=(1, 3, 2), fillstyle=AmscoFillStyle.continuous)
 886     'hereissometexttoencipher'
 887     """
 888
 889     grid = amsco_transposition_positions(message, keyword,
 890         fillpattern=fillpattern, fillstyle=fillstyle)
 891     transposed_sections = [s for c in [l for l in zip(*grid)] for s in c]
 892     plaintext_list = [''] * len(transposed_sections)
 893     current_pos = 0
 894     for slice in transposed_sections:
 895         plaintext_list[slice.index] = message[current_pos:current_pos-slice.start+slice.end][:len(message[slice.start:slice.end])]
 896         current_pos += len(message[slice.start:slice.end])
 897     return cat(plaintext_list)
 898
 899
 900 class PocketEnigma(object):
 901     """A pocket enigma machine
 902     The wheel is internally represented as a 26-element list self.wheel_map,
 903     where wheel_map[i] == j shows that the position i places on from the arrow
 904     maps to the position j places on.
 905     """
 906     def __init__(self, wheel=1, position='a'):
 907         """initialise the pocket enigma, including which wheel to use and the
 908         starting position of the wheel.
 909
 910         The wheel is either 1 or 2 (the predefined wheels) or a list of letter
 911         pairs.
 912
 913         The position is the letter pointed to by the arrow on the wheel.
 914
 915         >>> pe.wheel_map
 916         [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]
 917         >>> pe.position
 918         0
 919         """
 920         self.wheel1 = [('a', 'z'), ('b', 'e'), ('c', 'x'), ('d', 'k'),
 921             ('f', 'h'), ('g', 'j'), ('i', 'm'), ('l', 'r'), ('n', 'o'),
 922             ('p', 'v'), ('q', 't'), ('s', 'u'), ('w', 'y')]
 923         self.wheel2 = [('a', 'c'), ('b', 'd'), ('e', 'w'), ('f', 'i'),
 924             ('g', 'p'), ('h', 'm'), ('j', 'k'), ('l', 'n'), ('o', 'q'),
 925             ('r', 'z'), ('s', 'u'), ('t', 'v'), ('x', 'y')]
 926         if wheel == 1:
 927             self.make_wheel_map(self.wheel1)
 928         elif wheel == 2:
 929             self.make_wheel_map(self.wheel2)
 930         else:
 931             self.validate_wheel_spec(wheel)
 932             self.make_wheel_map(wheel)
 933         if position in string.ascii_lowercase:
 934             self.position = ord(position) - ord('a')
 935         else:
 936             self.position = position
 937
 938     def make_wheel_map(self, wheel_spec):
 939         """Expands a wheel specification from a list of letter-letter pairs
 940         into a full wheel_map.
 941
 942         >>> pe.make_wheel_map(pe.wheel2)
 943         [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]
 944         """
 945         self.validate_wheel_spec(wheel_spec)
 946         self.wheel_map = [0] * 26
 947         for p in wheel_spec:
 948             self.wheel_map[ord(p[0]) - ord('a')] = ord(p[1]) - ord('a')
 949             self.wheel_map[ord(p[1]) - ord('a')] = ord(p[0]) - ord('a')
 950         return self.wheel_map
 951
 952     def validate_wheel_spec(self, wheel_spec):
 953         """Validates that a wheel specificaiton will turn into a valid wheel
 954         map.
 955
 956         >>> pe.validate_wheel_spec([])
 957         Traceback (most recent call last):
 958             ...
 959         ValueError: Wheel specification has 0 pairs, requires 13
 960         >>> pe.validate_wheel_spec([('a', 'b', 'c')]*13)
 961         Traceback (most recent call last):
 962             ...
 963         ValueError: Not all mappings in wheel specificationhave two elements
 964         >>> pe.validate_wheel_spec([('a', 'b')]*13)
 965         Traceback (most recent call last):
 966             ...
 967         ValueError: Wheel specification does not contain 26 letters
 968         """
 969         if len(wheel_spec) != 13:
 970             raise ValueError("Wheel specification has {} pairs, requires 13".
 971                 format(len(wheel_spec)))
 972         for p in wheel_spec:
 973             if len(p) != 2:
 974                 raise ValueError("Not all mappings in wheel specification"
 975                     "have two elements")
 976         if len(set([p[0] for p in wheel_spec] +
 977                     [p[1] for p in wheel_spec])) != 26:
 978             raise ValueError("Wheel specification does not contain 26 letters")
 979
 980     def encipher_letter(self, letter):
 981         """Enciphers a single letter, by advancing the wheel before looking up
 982         the letter on the wheel.
 983
 984         >>> pe.set_position('f')
 985         5
 986         >>> pe.encipher_letter('k')
 987         'h'
 988         """
 989         self.advance()
 990         return self.lookup(letter)
 991     decipher_letter = encipher_letter
 992
 993     def lookup(self, letter):
 994         """Look up what a letter enciphers to, without turning the wheel.
 995
 996         >>> pe.set_position('f')
 997         5
 998         >>> cat([pe.lookup(l) for l in string.ascii_lowercase])
 999         'udhbfejcpgmokrliwntsayqzvx'
1000         >>> pe.lookup('A')
1001         ''
1002         """
1003         if letter in string.ascii_lowercase:
1004             return chr(
1005                 (self.wheel_map[(ord(letter) - ord('a') - self.position) % 26] +
1006                     self.position) % 26 +
1007                 ord('a'))
1008         else:
1009             return ''
1010
1011     def advance(self):
1012         """Advances the wheel one position.
1013
1014         >>> pe.set_position('f')
1015         5
1016         >>> pe.advance()
1017         6
1018         """
1019         self.position = (self.position + 1) % 26
1020         return self.position
1021
1022     def encipher(self, message, starting_position=None):
1023         """Enciphers a whole message.
1024
1025         >>> pe.set_position('f')
1026         5
1027         >>> pe.encipher('helloworld')
1028         'kjsglcjoqc'
1029         >>> pe.set_position('f')
1030         5
1031         >>> pe.encipher('kjsglcjoqc')
1032         'helloworld'
1033         >>> pe.encipher('helloworld', starting_position = 'x')
1034         'egrekthnnf'
1035         """
1036         if starting_position:
1037             self.set_position(starting_position)
1038         transformed = ''
1039         for l in message:
1040             transformed += self.encipher_letter(l)
1041         return transformed
1042     decipher = encipher
1043
1044     def set_position(self, position):
1045         """Sets the position of the wheel, by specifying the letter the arrow
1046         points to.
1047
1048         >>> pe.set_position('a')
1049         0
1050         >>> pe.set_position('m')
1051         12
1052         >>> pe.set_position('z')
1053         25
1054         """
1055         self.position = ord(position) - ord('a')
1056         return self.position
1057
1058
1059 if __name__ == "__main__":
1060     import doctest
1061     doctest.testmod(extraglobs={'pe': PocketEnigma(1, 'a')})