Worked on Enigma, mainly changing how the notch positions are handled
[cipher-training.git] / test_bombe.py
1 import unittest
2 import string
3
4 from enigma import *
5 from bombe import *
6
7 class ScramblerTest(unittest.TestCase):
8 def setUp(self):
9 self.scrambler = Scrambler(wheel_i_spec, wheel_ii_spec,
10 wheel_iii_spec, reflector_b_spec)
11
12 def test_attributes(self):
13 self.assertEqual(self.scrambler.wheel_positions, (0, 0, 0))
14 self.assertEqual(self.scrambler.wheel_positions_l, ('a', 'a', 'a'))
15
16 def test_set_positions(self):
17 self.scrambler.set_positions(1, 2, 3)
18 self.assertEqual(self.scrambler.wheel_positions, (1, 2, 3))
19 self.assertEqual(self.scrambler.wheel_positions_l, ('b', 'c', 'd'))
20 self.scrambler.set_positions('p', 'q', 'r')
21 self.assertEqual(self.scrambler.wheel_positions, (15, 16, 17))
22 self.assertEqual(self.scrambler.wheel_positions_l, ('p', 'q', 'r'))
23
24 def test_advance(self):
25 self.assertEqual(self.scrambler.wheel_positions, (0, 0, 0))
26 self.scrambler.advance()
27 self.assertEqual(self.scrambler.wheel_positions, (0, 0, 1))
28 self.scrambler.advance()
29 self.assertEqual(self.scrambler.wheel_positions, (0, 0, 2))
30 self.scrambler.set_positions(0, 0, 25)
31 self.assertEqual(self.scrambler.wheel_positions, (0, 0, 25))
32 self.scrambler.advance()
33 self.assertEqual(self.scrambler.wheel_positions, (0, 0, 0))
34 self.scrambler.set_positions(0, 0, 25)
35 self.scrambler.advance(wheel3=False)
36 self.assertEqual(self.scrambler.wheel_positions, (0, 0, 25))
37 self.scrambler.set_positions(0, 0, 25)
38 self.scrambler.advance(wheel2=True)
39 self.assertEqual(self.scrambler.wheel_positions, (0, 1, 0))
40 self.scrambler.set_positions(0, 0, 25)
41 self.scrambler.advance(wheel1=True, wheel2=True)
42 self.assertEqual(self.scrambler.wheel_positions, (1, 1, 0))
43
44 def test_lookups(self):
45 self.scrambler.set_positions(0, 0, 0)
46 self.assertEqual(cat(self.scrambler.lookup(l)
47 for l in string.ascii_lowercase),
48 'uejobtpzwcnsrkdgvmlfaqiyxh')
49 self.assertEqual(cat(self.scrambler.lookup(l)
50 for l in 'uejobtpzwcnsrkdgvmlfaqiyxh'),
51 'abcdefghijklmnopqrstuvwxyz')
52 self.scrambler.set_positions('p', 'q', 'r')
53 self.assertEqual(cat(self.scrambler.lookup(l)
54 for l in string.ascii_lowercase),
55 'jgqmnwbtvaurdezxclyhkifpso')
56 self.assertEqual(cat(self.scrambler.lookup(l)
57 for l in 'jgqmnwbtvaurdezxclyhkifpso'),
58 'abcdefghijklmnopqrstuvwxyz')
59
60 class BombeTest(unittest.TestCase):
61 def setUp(self):
62 self.bombe = Bombe(wheel_i_spec, wheel_ii_spec,
63 wheel_iii_spec, reflector_b_spec)
64 self.plaintext = 'thisisatestmessage'
65 self.ciphertext = 'opgndxcrwomnlnecjz'
66 self.menu = make_menu(self.plaintext, self.ciphertext)
67 self.bombe.read_menu(self.menu)
68
69 def test_menu(self):
70 self.assertEqual(len(self.bombe.connections), 18)
71 self.assertEqual(':'.join(sorted(cat(sorted(c.banks))
72 for c in self.bombe.connections)),
73 'ac:ac:di:el:es:ew:ez:gi:gj:hp:mn:mt:ns:ns:os:ot:rt:sx')
74 self.assertEqual(':'.join(sorted(cat(c.scrambler.wheel_positions_l)
75 for c in self.bombe.connections)),
76 'aaa:aab:aac:aad:aae:aaf:aag:aah:aai:aaj:aak:aal:aam:aan:aao:aap:aaq:aar')
77
78 self.bombe.read_menu(self.menu)
79 self.assertEqual(len(self.bombe.connections), 18)
80
81 def test_signal(self):
82 self.bombe.test(Signal('t', 't'))
83 self.assertEqual(len(self.bombe.banks['t']), 26)
84 self.assertTrue(all(self.bombe.banks['t'].values()))
85 self.assertEqual(sum(1 for s in self.bombe.banks['u'].values() if s), 18)
86
87 self.bombe.set_positions('a', 'a', 'b')
88 self.bombe.test()
89 self.assertEqual(sum(1 for b in self.bombe.banks
90 for s in self.bombe.banks[b].values() if s),
91 11)
92
93 def test_valid_with_rings(self):
94 pt31 = 'someplaintext'
95 ct31 = 'dhnpforeeimgg'
96 menu31 = make_menu(pt31, ct31)
97 b31 = Bombe(wheel_i_spec, wheel_v_spec, wheel_iii_spec, reflector_b_spec)
98 b31.read_menu(menu31)
99 b31.set_positions('e', 'l', 'f')
100
101 b31.test(Signal('s', 'o'))
102 self.assertEqual(sum(1 for b in b31.banks
103 for s in b31.banks[b].values() if s),
104 5)
105 self.assertEqual(':'.join(sorted(cat(sorted(p))
106 for p in b31.possible_plugboards())),
107 'd:hl:os')
108
109 b31.test(Signal('o', 'o'))
110 self.assertEqual(sum(1 for b in b31.banks
111 for s in b31.banks[b].values() if s),
112 507)
113 self.assertEqual(':'.join(sorted(cat(sorted(p))
114 for p in b31.possible_plugboards())),
115 'bg:ey:fp:in:m:tx')
116
117 def test_invalid_with_rings(self):
118 pt31 = 'someplaintext'
119 ct31 = 'dhnpforeeimgg'
120 menu31 = make_menu(pt31, ct31)
121 b31 = Bombe(wheel_i_spec, wheel_v_spec, wheel_iii_spec, reflector_b_spec)
122 b31.read_menu(menu31)
123 b31.set_positions('a', 'a', 'a')
124
125 b31.test(Signal('a', 'o'))
126 self.assertEqual(sum(1 for b in b31.banks
127 for s in b31.banks[b].values() if s),
128 514)
129 self.assertEqual(':'.join(sorted(cat(sorted(p))
130 for p in b31.possible_plugboards())),
131 '')
132
133 if __name__ == '__main__':
134 unittest.main()