szyfrow/polybius.py

   1 """Simple digraph substitution cipher, using the
   2 [Polybius square](https://en.wikipedia.org/wiki/Polybius_square). Enciphering
   3 and deciphering, and a couple of ways to break these ciphers.
   4 """
   5 import multiprocessing
   6 from szyfrow.support.utilities import *
   7 from szyfrow.support.language_models import *
   8 from szyfrow.keyword_cipher import KeywordWrapAlphabet, keyword_cipher_alphabet_of
   9
  10 def polybius_grid(keyword, column_order, row_order, letters_to_merge=None,
  11                   wrap_alphabet=KeywordWrapAlphabet.from_a):
  12     """Grid for a Polybius cipher, using a keyword to rearrange the
  13     alphabet.
  14
  15
  16     >>> polybius_grid('a', 'abcde', 'abcde')['x'] == ('e', 'c')
  17     True
  18     >>> polybius_grid('elephant', 'abcde', 'abcde')['e'] == ('a', 'a')
  19     True
  20     >>> polybius_grid('elephant', 'abcde', 'abcde')['b'] == ('b', 'c')
  21     True
  22     """
  23     alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet=wrap_alphabet)
  24     if letters_to_merge is None:
  25         letters_to_merge = {'j': 'i'}
  26     grid = {l: k
  27             for k, l in zip([(c, r) for c in column_order for r in row_order],
  28                 [l for l in alphabet if l not in letters_to_merge])}
  29     for l in letters_to_merge:
  30         grid[l] = grid[letters_to_merge[l]]
  31     return grid
  32
  33 def polybius_reverse_grid(keyword, column_order, row_order, letters_to_merge=None,
  34                   wrap_alphabet=KeywordWrapAlphabet.from_a):
  35     """Grid for decrypting using a Polybius cipher, using a keyword to
  36     rearrange the alphabet.
  37
  38     >>> polybius_reverse_grid('a', 'abcde', 'abcde')['e', 'c'] == 'x'
  39     True
  40     >>> polybius_reverse_grid('elephant', 'abcde', 'abcde')['a', 'a'] == 'e'
  41     True
  42     >>> polybius_reverse_grid('elephant', 'abcde', 'abcde')['b', 'c'] == 'b'
  43     True
  44     """
  45     alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet=wrap_alphabet)
  46     if letters_to_merge is None:
  47         letters_to_merge = {'j': 'i'}
  48     grid = {k: l
  49             for k, l in zip([(c, r) for c in column_order for r in row_order],
  50                 [l for l in alphabet if l not in letters_to_merge])}
  51     return grid
  52
  53
  54 def polybius_flatten(pair, column_first):
  55     """Convert a pair of characters into a single string."""
  56     if column_first:
  57         return str(pair[1]) + str(pair[0])
  58     else:
  59         return str(pair[0]) + str(pair[1])
  60
  61 def polybius_encipher(message, keyword, column_order, row_order,
  62                       column_first=False,
  63                       letters_to_merge=None, wrap_alphabet=KeywordWrapAlphabet.from_a):
  64     """Encipher a message with Polybius cipher, using a keyword to rearrange
  65     the alphabet
  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]))
 108           for p in chunks(message, 2)]
 109     else:
 110         pairs = [(row_index_type(p[0]), column_index_type(p[1]))
 111           for p in chunks(message, 2)]
 112     return cat(grid[p] for p in pairs if p in grid)
 113
 114
 115 def polybius_break(message, column_labels, row_labels,
 116                       letters_to_merge=None,
 117                       wordlist=keywords, fitness=Pletters,
 118                       number_of_solutions=1, chunksize=500):
 119     """Breaks a Polybius substitution cipher using a dictionary and
 120     frequency analysis
 121
 122     >>> polybius_break_mp(polybius_encipher('this is a test message for the ' \
 123           'polybius decipherment', 'elephant', 'abcde', 'abcde'), \
 124           'abcde', 'abcde', \
 125           wordlist=['cat', 'elephant', 'kangaroo']) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
 126     (('elephant', <KeywordWrapAlphabet.from_a: 1>, 'abcde', 'abcde', False), \
 127     -54.53880...)
 128     >>> polybius_break_mp(polybius_encipher('this is a test message for the ' \
 129           'polybius decipherment', 'elephant', 'abcde', 'abcde', column_first=True), \
 130           'abcde', 'abcde', \
 131           wordlist=['cat', 'elephant', 'kangaroo']) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
 132     (('elephant', <KeywordWrapAlphabet.from_a: 1>, 'abcde', 'abcde', True), \
 133     -54.53880...)
 134     >>> polybius_break_mp(polybius_encipher('this is a test message for the ' \
 135           'polybius decipherment', 'elephant', 'abcde', 'abcde', column_first=False), \
 136           'abcde', 'abcde', \
 137           wordlist=['cat', 'elephant', 'kangaroo']) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
 138     (('elephant', <KeywordWrapAlphabet.from_a: 1>, 'abcde', 'abcde', False), \
 139     -54.53880...)
 140     >>> polybius_break_mp(polybius_encipher('this is a test message for the ' \
 141           'polybius decipherment', 'elephant', 'abcde', 'pqrst', column_first=True), \
 142           'abcde', 'pqrst', \
 143           wordlist=['cat', 'elephant', 'kangaroo']) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
 144     (('elephant', <KeywordWrapAlphabet.from_a: 1>, 'abcde', 'pqrst', True), \
 145     -54.53880...)
 146     """
 147     if letters_to_merge is None:
 148         letters_to_merge = {'j': 'i'}
 149     with multiprocessing.Pool() as pool:
 150         helper_args = [(message, word, wrap,
 151                         column_labels, row_labels, column_first,
 152                         letters_to_merge,
 153                         fitness)
 154                        for word in wordlist
 155                        for wrap in KeywordWrapAlphabet
 156                        for column_first in [False, True]]
 157         # Gotcha: the helper function here needs to be defined at the top level
 158         #   (limitation of Pool.starmap)
 159         breaks = pool.starmap(polybius_break_worker, helper_args, chunksize)
 160         if number_of_solutions == 1:
 161             return max(breaks, key=lambda k: k[1])
 162         else:
 163             return sorted(breaks, key=lambda k: k[1], reverse=True)[:number_of_solutions]
 164
 165 def polybius_break_worker(message, keyword, wrap_alphabet,
 166                           column_order, row_order, column_first,
 167                           letters_to_merge,
 168                           fitness):
 169     plaintext = polybius_decipher(message, keyword,
 170                                   column_order, row_order,
 171                                   column_first=column_first,
 172                                   letters_to_merge=letters_to_merge,
 173                                   wrap_alphabet=wrap_alphabet)
 174     if plaintext:
 175         fit = fitness(plaintext)
 176     else:
 177         fit = float('-inf')
 178     return (keyword, wrap_alphabet, column_order, row_order, column_first), fit
 179
 180 if __name__ == "__main__":
 181     import doctest