szyfrow/polybius.py

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