polybius.py

   1 from utilities import *
   2 from language_models import *
   3 from multiprocessing import Pool
   4
   5 from logger import logger
   6
   7 def polybius_grid(keyword, column_order, row_order, letters_to_merge=None,
   8                   wrap_alphabet=KeywordWrapAlphabet.from_a):
   9     """Grid for a Polybius cipher, using a keyword to rearrange the
  10     alphabet.
  11
  12
  13     >>> polybius_grid('a', 'abcde', 'abcde')['x'] == ('e', 'c')
  14     True
  15     >>> polybius_grid('elephant', 'abcde', 'abcde')['e'] == ('a', 'a')
  16     True
  17     >>> polybius_grid('elephant', 'abcde', 'abcde')['b'] == ('b', 'c')
  18     True
  19     """
  20     alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet=wrap_alphabet)
  21     if letters_to_merge is None:
  22         letters_to_merge = {'j': 'i'}
  23     grid = {l: k
  24             for k, l in zip([(c, r) for c in column_order for r in row_order],
  25                 [l for l in alphabet if l not in letters_to_merge])}
  26     for l in letters_to_merge:
  27         grid[l] = grid[letters_to_merge[l]]
  28     return grid
  29
  30 def polybius_reverse_grid(keyword, column_order, row_order, letters_to_merge=None,
  31                   wrap_alphabet=KeywordWrapAlphabet.from_a):
  32     """Grid for decrypting using a Polybius cipher, using a keyword to
  33     rearrange the alphabet.
  34
  35     >>> polybius_reverse_grid('a', 'abcde', 'abcde')['e', 'c'] == 'x'
  36     True
  37     >>> polybius_reverse_grid('elephant', 'abcde', 'abcde')['a', 'a'] == 'e'
  38     True
  39     >>> polybius_reverse_grid('elephant', 'abcde', 'abcde')['b', 'c'] == 'b'
  40     True
  41     """
  42     alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet=wrap_alphabet)
  43     if letters_to_merge is None:
  44         letters_to_merge = {'j': 'i'}
  45     grid = {k: l
  46             for k, l in zip([(c, r) for c in column_order for r in row_order],
  47                 [l for l in alphabet if l not in letters_to_merge])}
  48     return grid
  49
  50
  51 def polybius_flatten(pair, column_first):
  52     """Convert a series of pairs into a single list of characters"""
  53     if column_first:
  54         return str(pair[1]) + str(pair[0])
  55     else:
  56         return str(pair[0]) + str(pair[1])
  57
  58 def polybius_encipher(message, keyword, column_order, row_order,
  59                       column_first=False,
  60                       letters_to_merge=None, wrap_alphabet=KeywordWrapAlphabet.from_a):
  61     """Encipher a message with Polybius cipher, using a keyword to rearrange
  62     the alphabet
  63
  64
  65     >>> polybius_encipher('this is a test message for the ' \
  66           'polybius decipherment', 'elephant', \
  67           [1, 2, 3, 4, 5], [1, 2, 3, 4, 5], \
  68           wrap_alphabet=KeywordWrapAlphabet.from_last)
  69     '2214445544551522115522511155551543114252542214111352123234442355411135441314115451112122'
  70     >>> polybius_encipher('this is a test message for the ' \
  71           'polybius decipherment', 'elephant', 'abcde', 'abcde', \
  72           column_first=False)
  73     'bbadccddccddaebbaaddbbceaaddddaecbaacadadcbbadaaacdaabedbcccdeddbeaabdccacadaadcceaababb'
  74     >>> polybius_encipher('this is a test message for the ' \
  75           'polybius decipherment', 'elephant', 'abcde', 'abcde', \
  76           column_first=True)
  77     'bbdaccddccddeabbaaddbbecaaddddeabcaaacadcdbbdaaacaadbadecbccedddebaadbcccadaaacdecaaabbb'
  78     """
  79     grid = polybius_grid(keyword, column_order, row_order, letters_to_merge, wrap_alphabet)
  80     return cat(polybius_flatten(grid[l], column_first)
  81                for l in message
  82                if l in grid)
  83
  84
  85 def polybius_decipher(message, keyword, column_order, row_order,
  86                       column_first=False,
  87                       letters_to_merge=None, wrap_alphabet=KeywordWrapAlphabet.from_a):
  88     """Decipher a message with a Polybius cipher, using a keyword to rearrange
  89     the alphabet
  90
  91     >>> polybius_decipher('bbdaccddccddeabbaaddbbecaaddddeabcaaacadcdbbdaaaca'\
  92     'adbadecbccedddebaadbcccadaaacdecaaabbb', 'elephant', 'abcde', 'abcde', \
  93     column_first=False)
  94     'toisisvtestxessvbephktoefhnugiysweqifoekxelt'
  95
  96     >>> polybius_decipher('bbdaccddccddeabbaaddbbecaaddddeabcaaacadcdbbdaaaca'\
  97     'adbadecbccedddebaadbcccadaaacdecaaabbb', 'elephant', 'abcde', 'abcde', \
  98     column_first=True)
  99     'thisisatestmessageforthepolybiusdecipherment'
 100     """
 101     grid = polybius_reverse_grid(keyword, column_order, row_order, letters_to_merge, wrap_alphabet)
 102     column_index_type = type(column_order[0])
 103     row_index_type = type(row_order[0])
 104     if column_first:
 105         pairs = [(column_index_type(p[1]), row_index_type(p[0])) for p in chunks(message, 2)]
 106     else:
 107         pairs = [(row_index_type(p[0]), column_index_type(p[1])) for p in chunks(message, 2)]
 108     return cat(grid[p] for p in pairs if p in grid)
 109
 110
 111 def polybius_break_mp(message, column_labels, row_labels,
 112                       letters_to_merge=None,
 113                       wordlist=keywords, fitness=Pletters,
 114                       number_of_solutions=1, chunksize=500):
 115     """Breaks a Polybius substitution cipher using a dictionary and
 116     frequency analysis
 117
 118     >>> polybius_break_mp(polybius_encipher('this is a test message for the ' \
 119           'polybius decipherment', 'elephant', 'abcde', 'abcde'), \
 120           'abcde', 'abcde', \
 121           wordlist=['cat', 'elephant', 'kangaroo']) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
 122     (('elephant', <KeywordWrapAlphabet.from_a: 1>, 'abcde', 'abcde', False), \
 123     -54.53880...)
 124     >>> polybius_break_mp(polybius_encipher('this is a test message for the ' \
 125           'polybius decipherment', 'elephant', 'abcde', 'abcde', column_first=True), \
 126           'abcde', 'abcde', \
 127           wordlist=['cat', 'elephant', 'kangaroo']) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
 128     (('elephant', <KeywordWrapAlphabet.from_a: 1>, 'abcde', 'abcde', True), \
 129     -54.53880...)
 130     >>> polybius_break_mp(polybius_encipher('this is a test message for the ' \
 131           'polybius decipherment', 'elephant', 'abcde', 'abcde', column_first=False), \
 132           'abcde', 'abcde', \
 133           wordlist=['cat', 'elephant', 'kangaroo']) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
 134     (('elephant', <KeywordWrapAlphabet.from_a: 1>, 'abcde', 'abcde', False), \
 135     -54.53880...)
 136     >>> polybius_break_mp(polybius_encipher('this is a test message for the ' \
 137           'polybius decipherment', 'elephant', 'abcde', 'pqrst', column_first=True), \
 138           'abcde', 'pqrst', \
 139           wordlist=['cat', 'elephant', 'kangaroo']) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
 140     (('elephant', <KeywordWrapAlphabet.from_a: 1>, 'abcde', 'pqrst', True), \
 141     -54.53880...)
 142     """
 143     if letters_to_merge is None:
 144         letters_to_merge = {'j': 'i'}
 145     with Pool() as pool:
 146         helper_args = [(message, word, wrap,
 147                         column_labels, row_labels, column_first,
 148                         letters_to_merge,
 149                         fitness)
 150                        for word in wordlist
 151                        for wrap in KeywordWrapAlphabet
 152                        for column_first in [False, True]]
 153         # Gotcha: the helper function here needs to be defined at the top level
 154         #   (limitation of Pool.starmap)
 155         breaks = pool.starmap(polybius_break_worker, helper_args, chunksize)
 156         if number_of_solutions == 1:
 157             return max(breaks, key=lambda k: k[1])
 158         else:
 159             return sorted(breaks, key=lambda k: k[1], reverse=True)[:number_of_solutions]
 160
 161 def polybius_break_worker(message, keyword, wrap_alphabet,
 162                           column_order, row_order, column_first,
 163                           letters_to_merge,
 164                           fitness):
 165     plaintext = polybius_decipher(message, keyword,
 166                                   column_order, row_order,
 167                                   column_first=column_first,
 168                                   letters_to_merge=letters_to_merge,
 169                                   wrap_alphabet=wrap_alphabet)
 170     if plaintext:
 171         fit = fitness(plaintext)
 172     else:
 173         fit = float('-inf')
 174     logger.debug('Polybius break attempt using key {0} (wrap={1}, merging {2}), '
 175                  'columns as {3}, rows as {4} (column_first={5}) '
 176                  'gives fit of {6} and decrypt starting: '
 177                  '{7}'.format(keyword, wrap_alphabet, letters_to_merge,
 178                               column_order, row_order, column_first,
 179                               fit, sanitise(plaintext)[:50]))
 180     return (keyword, wrap_alphabet, column_order, row_order, column_first), fit
 181