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()