cipher/polybius.py

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