cipher/polybius.py

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