Tidyied paths and things
[szyfrow.git] / szyfrow / bifid.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 bifid_grid(keyword, wrap_alphabet, letter_mapping):
7 """Create the grids for a Bifid cipher
8 """
9 cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
10 if letter_mapping is None:
11 letter_mapping = {'j': 'i'}
12 translation = ''.maketrans(letter_mapping)
13 cipher_alphabet = cat(collections.OrderedDict.fromkeys(cipher_alphabet.translate(translation)))
14 f_grid = {k: ((i // 5) + 1, (i % 5) + 1)
15 for i, k in enumerate(cipher_alphabet)}
16 r_grid = {((i // 5) + 1, (i % 5) + 1): k
17 for i, k in enumerate(cipher_alphabet)}
18 return translation, f_grid, r_grid
19
20 def bifid_encipher(message, keyword, wrap_alphabet=KeywordWrapAlphabet.from_a,
21 letter_mapping=None, period=None, fillvalue=None):
22 """Bifid cipher
23
24 >>> bifid_encipher("indiajelly", 'iguana')
25 'ibidonhprm'
26 >>> bifid_encipher("indiacurry", 'iguana', period=4)
27 'ibnhgaqltm'
28 >>> bifid_encipher("indiacurry", 'iguana', period=4, fillvalue='x')
29 'ibnhgaqltzml'
30 """
31 translation, f_grid, r_grid = bifid_grid(keyword, wrap_alphabet, letter_mapping)
32
33 t_message = message.translate(translation)
34 pairs0 = [f_grid[l] for l in sanitise(t_message)]
35 if period:
36 chunked_pairs = [pairs0[i:i+period] for i in range(0, len(pairs0), period)]
37 if len(chunked_pairs[-1]) < period and fillvalue:
38 chunked_pairs[-1] += [f_grid[fillvalue]] * (period - len(chunked_pairs[-1]))
39 else:
40 chunked_pairs = [pairs0]
41
42 pairs1 = []
43 for c in chunked_pairs:
44 items = sum(list(list(i) for i in zip(*c)), [])
45 p = [(items[i], items[i+1]) for i in range(0, len(items), 2)]
46 pairs1 += p
47
48 return cat(r_grid[p] for p in pairs1)
49
50
51 def bifid_decipher(message, keyword, wrap_alphabet=KeywordWrapAlphabet.from_a,
52 letter_mapping=None, period=None, fillvalue=None):
53 """Decipher with bifid cipher
54
55 >>> bifid_decipher('ibidonhprm', 'iguana')
56 'indiaielly'
57 >>> bifid_decipher("ibnhgaqltm", 'iguana', period=4)
58 'indiacurry'
59 >>> bifid_decipher("ibnhgaqltzml", 'iguana', period=4)
60 'indiacurryxx'
61 """
62 translation, f_grid, r_grid = bifid_grid(keyword, wrap_alphabet, letter_mapping)
63
64 t_message = message.translate(translation)
65 pairs0 = [f_grid[l] for l in sanitise(t_message)]
66 if period:
67 chunked_pairs = [pairs0[i:i+period] for i in range(0, len(pairs0), period)]
68 if len(chunked_pairs[-1]) < period and fillvalue:
69 chunked_pairs[-1] += [f_grid[fillvalue]] * (period - len(chunked_pairs[-1]))
70 else:
71 chunked_pairs = [pairs0]
72
73 pairs1 = []
74 for c in chunked_pairs:
75 items = [j for i in c for j in i]
76 gap = len(c)
77 p = [(items[i], items[i+gap]) for i in range(gap)]
78 pairs1 += p
79
80 return cat(r_grid[p] for p in pairs1)
81
82
83 def bifid_break_mp(message, wordlist=keywords, fitness=Pletters, max_period=10,
84 number_of_solutions=1, chunksize=500):
85 """Breaks a keyword substitution cipher using a dictionary and
86 frequency analysis
87
88 >>> bifid_break_mp(bifid_encipher('this is a test message for the ' \
89 'keyword decipherment', 'elephant', wrap_alphabet=KeywordWrapAlphabet.from_last), \
90 wordlist=['cat', 'elephant', 'kangaroo']) # doctest: +ELLIPSIS
91 (('elephant', <KeywordWrapAlphabet.from_last: 2>, 0), -52.834575011...)
92 >>> bifid_break_mp(bifid_encipher('this is a test message for the ' \
93 'keyword decipherment', 'elephant', wrap_alphabet=KeywordWrapAlphabet.from_last), \
94 wordlist=['cat', 'elephant', 'kangaroo'], \
95 number_of_solutions=2) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
96 [(('elephant', <KeywordWrapAlphabet.from_last: 2>, 0), -52.834575011...),
97 (('elephant', <KeywordWrapAlphabet.from_largest: 3>, 0), -52.834575011...)]
98 """
99 with multiprocessing.Pool() as pool:
100 helper_args = [(message, word, wrap, period, fitness)
101 for word in wordlist
102 for wrap in KeywordWrapAlphabet
103 for period in range(max_period+1)]
104 # Gotcha: the helper function here needs to be defined at the top level
105 # (limitation of Pool.starmap)
106 breaks = pool.starmap(bifid_break_worker, helper_args, chunksize)
107 if number_of_solutions == 1:
108 return max(breaks, key=lambda k: k[1])
109 else:
110 return sorted(breaks, key=lambda k: k[1], reverse=True)[:number_of_solutions]
111
112 def bifid_break_worker(message, keyword, wrap_alphabet, period, fitness):
113 plaintext = bifid_decipher(message, keyword, wrap_alphabet, period=period)
114 fit = fitness(plaintext)
115 return (keyword, wrap_alphabet, period), fit
116
117 if __name__ == "__main__":
118 import doctest