tests/test_bombe.py

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