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), '*')
 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 class PocketEnigma(object):
 511     """A pocket enigma machine
 512     The wheel is internally represented as a 26-element list self.wheel_map,
 513     where wheel_map[i] == j shows that the position i places on from the arrow
 514     maps to the position j places on.
 515     """
 516     def __init__(self, wheel=1, position='a'):
 517         """initialise the pocket enigma, including which wheel to use and the
 518         starting position of the wheel.
 519
 520         The wheel is either 1 or 2 (the predefined wheels) or a list of letter
 521         pairs.
 522
 523         The position is the letter pointed to by the arrow on the wheel.
 524
 525         >>> pe.wheel_map
 526         [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]
 527         >>> pe.position
 528         0
 529         """
 530         self.wheel1 = [('a', 'z'), ('b', 'e'), ('c', 'x'), ('d', 'k'),
 531             ('f', 'h'), ('g', 'j'), ('i', 'm'), ('l', 'r'), ('n', 'o'),
 532             ('p', 'v'), ('q', 't'), ('s', 'u'), ('w', 'y')]
 533         self.wheel2 = [('a', 'c'), ('b', 'd'), ('e', 'w'), ('f', 'i'),
 534             ('g', 'p'), ('h', 'm'), ('j', 'k'), ('l', 'n'), ('o', 'q'),
 535             ('r', 'z'), ('s', 'u'), ('t', 'v'), ('x', 'y')]
 536         if wheel == 1:
 537             self.make_wheel_map(self.wheel1)
 538         elif wheel == 2:
 539             self.make_wheel_map(self.wheel2)
 540         else:
 541             self.validate_wheel_spec(wheel)
 542             self.make_wheel_map(wheel)
 543         if position in string.ascii_lowercase:
 544             self.position = ord(position) - ord('a')
 545         else:
 546             self.position = position
 547
 548     def make_wheel_map(self, wheel_spec):
 549         """Expands a wheel specification from a list of letter-letter pairs
 550         into a full wheel_map.
 551
 552         >>> pe.make_wheel_map(pe.wheel2)
 553         [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]
 554         """
 555         self.validate_wheel_spec(wheel_spec)
 556         self.wheel_map = [0] * 26
 557         for p in wheel_spec:
 558             self.wheel_map[ord(p[0]) - ord('a')] = ord(p[1]) - ord('a')
 559             self.wheel_map[ord(p[1]) - ord('a')] = ord(p[0]) - ord('a')
 560         return self.wheel_map
 561
 562     def validate_wheel_spec(self, wheel_spec):
 563         """Validates that a wheel specificaiton will turn into a valid wheel
 564         map.
 565
 566         >>> pe.validate_wheel_spec([])
 567         Traceback (most recent call last):
 568             ...
 569         ValueError: Wheel specification has 0 pairs, requires 13
 570         >>> pe.validate_wheel_spec([('a', 'b', 'c')]*13)
 571         Traceback (most recent call last):
 572             ...
 573         ValueError: Not all mappings in wheel specificationhave two elements
 574         >>> pe.validate_wheel_spec([('a', 'b')]*13)
 575         Traceback (most recent call last):
 576             ...
 577         ValueError: Wheel specification does not contain 26 letters
 578         """
 579         if len(wheel_spec) != 13:
 580             raise ValueError("Wheel specification has {} pairs, requires 13".
 581                 format(len(wheel_spec)))
 582         for p in wheel_spec:
 583             if len(p) != 2:
 584                 raise ValueError("Not all mappings in wheel specification"
 585                     "have two elements")
 586         if len(set([p[0] for p in wheel_spec] +
 587                     [p[1] for p in wheel_spec])) != 26:
 588             raise ValueError("Wheel specification does not contain 26 letters")
 589
 590     def encipher_letter(self, letter):
 591         """Enciphers a single letter, by advancing the wheel before looking up
 592         the letter on the wheel.
 593
 594         >>> pe.set_position('f')
 595         5
 596         >>> pe.encipher_letter('k')
 597         'h'
 598         """
 599         self.advance()
 600         return self.lookup(letter)
 601     decipher_letter = encipher_letter
 602
 603     def lookup(self, letter):
 604         """Look up what a letter enciphers to, without turning the wheel.
 605
 606         >>> pe.set_position('f')
 607         5
 608         >>> ''.join([pe.lookup(l) for l in string.ascii_lowercase])
 609         'udhbfejcpgmokrliwntsayqzvx'
 610         >>> pe.lookup('A')
 611         ''
 612         """
 613         if letter in string.ascii_lowercase:
 614             return chr(
 615                 (self.wheel_map[(ord(letter) - ord('a') - self.position) % 26] +
 616                     self.position) % 26 +
 617                 ord('a'))
 618         else:
 619             return ''
 620
 621     def advance(self):
 622         """Advances the wheel one position.
 623
 624         >>> pe.set_position('f')
 625         5
 626         >>> pe.advance()
 627         6
 628         """
 629         self.position = (self.position + 1) % 26
 630         return self.position
 631
 632     def encipher(self, message, starting_position=None):
 633         """Enciphers a whole message.
 634
 635         >>> pe.set_position('f')
 636         5
 637         >>> pe.encipher('helloworld')
 638         'kjsglcjoqc'
 639         >>> pe.set_position('f')
 640         5
 641         >>> pe.encipher('kjsglcjoqc')
 642         'helloworld'
 643         >>> pe.encipher('helloworld', starting_position = 'x')
 644         'egrekthnnf'
 645         """
 646         if starting_position:
 647             self.set_position(starting_position)
 648         transformed = ''
 649         for l in message:
 650             transformed += self.encipher_letter(l)
 651         return transformed
 652     decipher = encipher
 653
 654     def set_position(self, position):
 655         """Sets the position of the wheel, by specifying the letter the arrow
 656         points to.
 657
 658         >>> pe.set_position('a')
 659         0
 660         >>> pe.set_position('m')
 661         12
 662         >>> pe.set_position('z')
 663         25
 664         """
 665         self.position = ord(position) - ord('a')
 666         return self.position
 667
 668
 669 if __name__ == "__main__":
 670     import doctest
 671     doctest.testmod(extraglobs={'pe': PocketEnigma(1, 'a')})