cipher.py

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