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                             )
 230         if one_based: plaintext_number -= 1
 231         return chr(plaintext_number % 26 + alphabet_start)
 232     else:
 233         return letter
 234
 235 def affine_encipher(message, multiplier=1, adder=0, one_based=True):
 236     """Encipher a message
 237
 238     >>> affine_encipher('hours passed during which jerico tried every ' \
 239            'trick he could think of', 15, 22, True)
 240     'lmyfu bkuusd dyfaxw claol psfaom jfasd snsfg jfaoe ls omytd jlaxe mh'
 241     """
 242     enciphered = [affine_encipher_letter(l, multiplier, adder, one_based)
 243                   for l in message]
 244     return ''.join(enciphered)
 245
 246 def affine_decipher(message, multiplier=1, adder=0, one_based=True):
 247     """Decipher a message
 248
 249     >>> affine_decipher('lmyfu bkuusd dyfaxw claol psfaom jfasd snsfg ' \
 250            'jfaoe ls omytd jlaxe mh', 15, 22, True)
 251     'hours passed during which jerico tried every trick he could think of'
 252     """
 253     enciphered = [affine_decipher_letter(l, multiplier, adder, one_based)
 254                   for l in message]
 255     return ''.join(enciphered)
 256
 257
 258 class KeywordWrapAlphabet(Enum):
 259     """Ways of wrapping the alphabet for keyword-based substitution ciphers."""
 260     from_a = 1
 261     from_last = 2
 262     from_largest = 3
 263
 264
 265 def keyword_cipher_alphabet_of(keyword,
 266         wrap_alphabet=KeywordWrapAlphabet.from_a):
 267     """Find the cipher alphabet given a keyword.
 268     wrap_alphabet controls how the rest of the alphabet is added
 269     after the keyword.
 270
 271     >>> keyword_cipher_alphabet_of('bayes')
 272     'bayescdfghijklmnopqrtuvwxz'
 273     >>> keyword_cipher_alphabet_of('bayes', KeywordWrapAlphabet.from_a)
 274     'bayescdfghijklmnopqrtuvwxz'
 275     >>> keyword_cipher_alphabet_of('bayes', KeywordWrapAlphabet.from_last)
 276     'bayestuvwxzcdfghijklmnopqr'
 277     >>> keyword_cipher_alphabet_of('bayes', KeywordWrapAlphabet.from_largest)
 278     'bayeszcdfghijklmnopqrtuvwx'
 279     """
 280     if wrap_alphabet == KeywordWrapAlphabet.from_a:
 281         cipher_alphabet = ''.join(deduplicate(sanitise(keyword) +
 282                                               string.ascii_lowercase))
 283     else:
 284         if wrap_alphabet == KeywordWrapAlphabet.from_last:
 285             last_keyword_letter = deduplicate(sanitise(keyword))[-1]
 286         else:
 287             last_keyword_letter = sorted(sanitise(keyword))[-1]
 288         last_keyword_position = string.ascii_lowercase.find(
 289             last_keyword_letter) + 1
 290         cipher_alphabet = ''.join(
 291             deduplicate(sanitise(keyword) +
 292                         string.ascii_lowercase[last_keyword_position:] +
 293                         string.ascii_lowercase))
 294     return cipher_alphabet
 295
 296
 297 def keyword_encipher(message, keyword,
 298                      wrap_alphabet=KeywordWrapAlphabet.from_a):
 299     """Enciphers a message with a keyword substitution cipher.
 300     wrap_alphabet controls how the rest of the alphabet is added
 301     after the keyword.
 302     0 : from 'a'
 303     1 : from the last letter in the sanitised keyword
 304     2 : from the largest letter in the sanitised keyword
 305
 306     >>> keyword_encipher('test message', 'bayes')
 307     'rsqr ksqqbds'
 308     >>> keyword_encipher('test message', 'bayes', KeywordWrapAlphabet.from_a)
 309     'rsqr ksqqbds'
 310     >>> keyword_encipher('test message', 'bayes', KeywordWrapAlphabet.from_last)
 311     'lskl dskkbus'
 312     >>> keyword_encipher('test message', 'bayes', KeywordWrapAlphabet.from_largest)
 313     'qspq jsppbcs'
 314     """
 315     cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
 316     cipher_translation = ''.maketrans(string.ascii_lowercase, cipher_alphabet)
 317     return unaccent(message).lower().translate(cipher_translation)
 318
 319 def keyword_decipher(message, keyword,
 320                      wrap_alphabet=KeywordWrapAlphabet.from_a):
 321     """Deciphers a message with a keyword substitution cipher.
 322     wrap_alphabet controls how the rest of the alphabet is added
 323     after the keyword.
 324     0 : from 'a'
 325     1 : from the last letter in the sanitised keyword
 326     2 : from the largest letter in the sanitised keyword
 327
 328     >>> keyword_decipher('rsqr ksqqbds', 'bayes')
 329     'test message'
 330     >>> keyword_decipher('rsqr ksqqbds', 'bayes', KeywordWrapAlphabet.from_a)
 331     'test message'
 332     >>> keyword_decipher('lskl dskkbus', 'bayes', KeywordWrapAlphabet.from_last)
 333     'test message'
 334     >>> keyword_decipher('qspq jsppbcs', 'bayes', KeywordWrapAlphabet.from_largest)
 335     'test message'
 336     """
 337     cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
 338     cipher_translation = ''.maketrans(cipher_alphabet, string.ascii_lowercase)
 339     return message.lower().translate(cipher_translation)
 340
 341
 342 def vigenere_encipher(message, keyword):
 343     """Vigenere encipher
 344
 345     >>> vigenere_encipher('hello', 'abc')
 346     'hfnlp'
 347     """
 348     shifts = [ord(l) - ord('a') for l in sanitise(keyword)]
 349     pairs = zip(message, cycle(shifts))
 350     return ''.join([caesar_encipher_letter(l, k) for l, k in pairs])
 351
 352 def vigenere_decipher(message, keyword):
 353     """Vigenere decipher
 354
 355     >>> vigenere_decipher('hfnlp', 'abc')
 356     'hello'
 357     """
 358     shifts = [ord(l) - ord('a') for l in sanitise(keyword)]
 359     pairs = zip(message, cycle(shifts))
 360     return ''.join([caesar_decipher_letter(l, k) for l, k in pairs])
 361
 362 beaufort_encipher = vigenere_decipher
 363 beaufort_decipher = vigenere_encipher
 364
 365
 366 def transpositions_of(keyword):
 367     """Finds the transpostions given by a keyword. For instance, the keyword
 368     'clever' rearranges to 'celrv', so the first column (0) stays first, the
 369     second column (1) moves to third, the third column (2) moves to second,
 370     and so on.
 371
 372     If passed a tuple, assume it's already a transposition and just return it.
 373
 374     >>> transpositions_of('clever')
 375     (0, 2, 1, 4, 3)
 376     >>> transpositions_of('fred')
 377     (3, 2, 0, 1)
 378     >>> transpositions_of((3, 2, 0, 1))
 379     (3, 2, 0, 1)
 380     """
 381     if isinstance(keyword, tuple):
 382         return keyword
 383     else:
 384         key = deduplicate(keyword)
 385         transpositions = tuple(key.index(l) for l in sorted(key))
 386         return transpositions
 387
 388 def pad(message_len, group_len, fillvalue):
 389     """Returns the padding required to extend a message of message_len to an
 390     even multiple of group_len, by adding repreated copies of fillvalue.
 391     fillvalue can either be a character or a function that returns a character.
 392
 393     >>> pad(10, 4, '!')
 394     '!!'
 395     >>> pad(8, 4, '!')
 396     ''
 397     >>> pad(16, 4, '!')
 398     ''
 399     >>> pad(10, 4, lambda: '*')
 400     '**'
 401     """
 402     padding_length = group_len - message_len % group_len
 403     if padding_length == group_len: padding_length = 0
 404     padding = ''
 405     for _ in range(padding_length):
 406         if callable(fillvalue):
 407             padding += fillvalue()
 408         else:
 409             padding += fillvalue
 410     return padding
 411
 412 def column_transposition_encipher(message, keyword, fillvalue=' ',
 413                                   fillcolumnwise=False,
 414                                   emptycolumnwise=False):
 415     """Enciphers using the column transposition cipher.
 416     Message is padded to allow all rows to be the same length.
 417
 418     >>> column_transposition_encipher('hellothere', 'abcdef', fillcolumnwise=True)
 419     'hlohr eltee '
 420     >>> column_transposition_encipher('hellothere', 'abcdef', fillcolumnwise=True, emptycolumnwise=True)
 421     'hellothere  '
 422     >>> column_transposition_encipher('hellothere', 'abcdef')
 423     'hellothere  '
 424     >>> column_transposition_encipher('hellothere', 'abcde')
 425     'hellothere'
 426     >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=True)
 427     'hellothere'
 428     >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=False)
 429     'hlohreltee'
 430     >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=True)
 431     'htehlelroe'
 432     >>> column_transposition_encipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=False)
 433     'hellothere'
 434     >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=True, emptycolumnwise=True)
 435     'heotllrehe'
 436     >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=True, emptycolumnwise=False)
 437     'holrhetlee'
 438     >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=False, emptycolumnwise=True)
 439     'htleehoelr'
 440     >>> column_transposition_encipher('hellothere', 'clever', fillcolumnwise=False, emptycolumnwise=False)
 441     'hleolteher'
 442     >>> column_transposition_encipher('hellothere', 'cleverly')
 443     'hleolthre e '
 444     >>> column_transposition_encipher('hellothere', 'cleverly', fillvalue='!')
 445     'hleolthre!e!'
 446     >>> column_transposition_encipher('hellothere', 'cleverly', fillvalue=lambda: '*')
 447     'hleolthre*e*'
 448     """
 449     transpositions = transpositions_of(keyword)
 450     message += pad(len(message), len(transpositions), fillvalue)
 451     if fillcolumnwise:
 452         rows = every_nth(message, len(message) // len(transpositions))
 453     else:
 454         rows = chunks(message, len(transpositions))
 455     transposed = [transpose(r, transpositions) for r in rows]
 456     if emptycolumnwise:
 457         return combine_every_nth(transposed)
 458     else:
 459         return ''.join(chain(*transposed))
 460
 461 def column_transposition_decipher(message, keyword, fillvalue=' ',
 462                                   fillcolumnwise=False,
 463                                   emptycolumnwise=False):
 464     """Deciphers using the column transposition cipher.
 465     Message is padded to allow all rows to be the same length.
 466
 467     >>> column_transposition_decipher('hellothere', 'abcde', fillcolumnwise=True, emptycolumnwise=True)
 468     'hellothere'
 469     >>> column_transposition_decipher('hlohreltee', 'abcde', fillcolumnwise=True, emptycolumnwise=False)
 470     'hellothere'
 471     >>> column_transposition_decipher('htehlelroe', 'abcde', fillcolumnwise=False, emptycolumnwise=True)
 472     'hellothere'
 473     >>> column_transposition_decipher('hellothere', 'abcde', fillcolumnwise=False, emptycolumnwise=False)
 474     'hellothere'
 475     >>> column_transposition_decipher('heotllrehe', 'clever', fillcolumnwise=True, emptycolumnwise=True)
 476     'hellothere'
 477     >>> column_transposition_decipher('holrhetlee', 'clever', fillcolumnwise=True, emptycolumnwise=False)
 478     'hellothere'
 479     >>> column_transposition_decipher('htleehoelr', 'clever', fillcolumnwise=False, emptycolumnwise=True)
 480     'hellothere'
 481     >>> column_transposition_decipher('hleolteher', 'clever', fillcolumnwise=False, emptycolumnwise=False)
 482     'hellothere'
 483     """
 484     transpositions = transpositions_of(keyword)
 485     message += pad(len(message), len(transpositions), '*')
 486     if emptycolumnwise:
 487         rows = every_nth(message, len(message) // len(transpositions))
 488     else:
 489         rows = chunks(message, len(transpositions))
 490     untransposed = [untranspose(r, transpositions) for r in rows]
 491     if fillcolumnwise:
 492         return combine_every_nth(untransposed)
 493     else:
 494         return ''.join(chain(*untransposed))
 495
 496 def scytale_encipher(message, rows, fillvalue=' '):
 497     """Enciphers using the scytale transposition cipher.
 498     Message is padded with spaces to allow all rows to be the same length.
 499
 500     >>> scytale_encipher('thequickbrownfox', 3)
 501     'tcnhkfeboqrxuo iw '
 502     >>> scytale_encipher('thequickbrownfox', 4)
 503     'tubnhirfecooqkwx'
 504     >>> scytale_encipher('thequickbrownfox', 5)
 505     'tubn hirf ecoo qkwx '
 506     >>> scytale_encipher('thequickbrownfox', 6)
 507     'tqcrnxhukof eibwo '
 508     >>> scytale_encipher('thequickbrownfox', 7)
 509     'tqcrnx hukof  eibwo  '
 510     """
 511     transpositions = [i for i in range(rows)]
 512     return column_transposition_encipher(message, transpositions,
 513             fillvalue=fillvalue, fillcolumnwise=True, emptycolumnwise=False)
 514
 515 def scytale_decipher(message, rows):
 516     """Deciphers using the scytale transposition cipher.
 517     Assumes the message is padded so that all rows are the same length.
 518
 519     >>> scytale_decipher('tcnhkfeboqrxuo iw ', 3)
 520     'thequickbrownfox  '
 521     >>> scytale_decipher('tubnhirfecooqkwx', 4)
 522     'thequickbrownfox'
 523     >>> scytale_decipher('tubn hirf ecoo qkwx ', 5)
 524     'thequickbrownfox    '
 525     >>> scytale_decipher('tqcrnxhukof eibwo ', 6)
 526     'thequickbrownfox  '
 527     >>> scytale_decipher('tqcrnx hukof  eibwo  ', 7)
 528     'thequickbrownfox     '
 529     """
 530     transpositions = [i for i in range(rows)]
 531     return column_transposition_decipher(message, transpositions,
 532         fillcolumnwise=True, emptycolumnwise=False)
 533
 534
 535 class PocketEnigma(object):
 536     """A pocket enigma machine
 537     The wheel is internally represented as a 26-element list self.wheel_map,
 538     where wheel_map[i] == j shows that the position i places on from the arrow
 539     maps to the position j places on.
 540     """
 541     def __init__(self, wheel=1, position='a'):
 542         """initialise the pocket enigma, including which wheel to use and the
 543         starting position of the wheel.
 544
 545         The wheel is either 1 or 2 (the predefined wheels) or a list of letter
 546         pairs.
 547
 548         The position is the letter pointed to by the arrow on the wheel.
 549
 550         >>> pe.wheel_map
 551         [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]
 552         >>> pe.position
 553         0
 554         """
 555         self.wheel1 = [('a', 'z'), ('b', 'e'), ('c', 'x'), ('d', 'k'),
 556             ('f', 'h'), ('g', 'j'), ('i', 'm'), ('l', 'r'), ('n', 'o'),
 557             ('p', 'v'), ('q', 't'), ('s', 'u'), ('w', 'y')]
 558         self.wheel2 = [('a', 'c'), ('b', 'd'), ('e', 'w'), ('f', 'i'),
 559             ('g', 'p'), ('h', 'm'), ('j', 'k'), ('l', 'n'), ('o', 'q'),
 560             ('r', 'z'), ('s', 'u'), ('t', 'v'), ('x', 'y')]
 561         if wheel == 1:
 562             self.make_wheel_map(self.wheel1)
 563         elif wheel == 2:
 564             self.make_wheel_map(self.wheel2)
 565         else:
 566             self.validate_wheel_spec(wheel)
 567             self.make_wheel_map(wheel)
 568         if position in string.ascii_lowercase:
 569             self.position = ord(position) - ord('a')
 570         else:
 571             self.position = position
 572
 573     def make_wheel_map(self, wheel_spec):
 574         """Expands a wheel specification from a list of letter-letter pairs
 575         into a full wheel_map.
 576
 577         >>> pe.make_wheel_map(pe.wheel2)
 578         [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]
 579         """
 580         self.validate_wheel_spec(wheel_spec)
 581         self.wheel_map = [0] * 26
 582         for p in wheel_spec:
 583             self.wheel_map[ord(p[0]) - ord('a')] = ord(p[1]) - ord('a')
 584             self.wheel_map[ord(p[1]) - ord('a')] = ord(p[0]) - ord('a')
 585         return self.wheel_map
 586
 587     def validate_wheel_spec(self, wheel_spec):
 588         """Validates that a wheel specificaiton will turn into a valid wheel
 589         map.
 590
 591         >>> pe.validate_wheel_spec([])
 592         Traceback (most recent call last):
 593             ...
 594         ValueError: Wheel specification has 0 pairs, requires 13
 595         >>> pe.validate_wheel_spec([('a', 'b', 'c')]*13)
 596         Traceback (most recent call last):
 597             ...
 598         ValueError: Not all mappings in wheel specificationhave two elements
 599         >>> pe.validate_wheel_spec([('a', 'b')]*13)
 600         Traceback (most recent call last):
 601             ...
 602         ValueError: Wheel specification does not contain 26 letters
 603         """
 604         if len(wheel_spec) != 13:
 605             raise ValueError("Wheel specification has {} pairs, requires"
 606                              " 13".format(len(wheel_spec)))
 607         for p in wheel_spec:
 608             if len(p) != 2:
 609                 raise ValueError("Not all mappings in wheel specification"
 610                                  "have two elements")
 611         if len(set([p[0] for p in wheel_spec] +
 612                    [p[1] for p in wheel_spec])) != 26:
 613             raise ValueError("Wheel specification does not contain 26 letters")
 614
 615     def encipher_letter(self, letter):
 616         """Enciphers a single letter, by advancing the wheel before looking up
 617         the letter on the wheel.
 618
 619         >>> pe.set_position('f')
 620         5
 621         >>> pe.encipher_letter('k')
 622         'h'
 623         """
 624         self.advance()
 625         return self.lookup(letter)
 626     decipher_letter = encipher_letter
 627
 628     def lookup(self, letter):
 629         """Look up what a letter enciphers to, without turning the wheel.
 630
 631         >>> pe.set_position('f')
 632         5
 633         >>> ''.join([pe.lookup(l) for l in string.ascii_lowercase])
 634         'udhbfejcpgmokrliwntsayqzvx'
 635         >>> pe.lookup('A')
 636         ''
 637         """
 638         if letter in string.ascii_lowercase:
 639             return chr(
 640                 (self.wheel_map[(ord(letter) - ord('a') - self.position) % 26] +
 641                     self.position) % 26 +
 642                 ord('a'))
 643         else:
 644             return ''
 645
 646     def advance(self):
 647         """Advances the wheel one position.
 648
 649         >>> pe.set_position('f')
 650         5
 651         >>> pe.advance()
 652         6
 653         """
 654         self.position = (self.position + 1) % 26
 655         return self.position
 656
 657     def encipher(self, message, starting_position=None):
 658         """Enciphers a whole message.
 659
 660         >>> pe.set_position('f')
 661         5
 662         >>> pe.encipher('helloworld')
 663         'kjsglcjoqc'
 664         >>> pe.set_position('f')
 665         5
 666         >>> pe.encipher('kjsglcjoqc')
 667         'helloworld'
 668         >>> pe.encipher('helloworld', starting_position = 'x')
 669         'egrekthnnf'
 670         """
 671         if starting_position:
 672             self.set_position(starting_position)
 673         transformed = ''
 674         for l in message:
 675             transformed += self.encipher_letter(l)
 676         return transformed
 677     decipher = encipher
 678
 679     def set_position(self, position):
 680         """Sets the position of the wheel, by specifying the letter the arrow
 681         points to.
 682
 683         >>> pe.set_position('a')
 684         0
 685         >>> pe.set_position('m')
 686         12
 687         >>> pe.set_position('z')
 688         25
 689         """
 690         self.position = ord(position) - ord('a')
 691         return self.position
 692
 693
 694 if __name__ == "__main__":
 695     import doctest
 696     doctest.testmod(extraglobs={'pe': PocketEnigma(1, 'a')})