cd78ee9f8a23445f9cc7355bac876f270b30a41f
[szyfrow.git] / 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