szyfrow/pocket_enigma.py

   1 """The Pocket Enigma machine, a simple example that illustrates the mechanisms
   2 of the Enigma machine. See [a review](http://www.savory.de/pocket_enigma.htm)
   3 for more information about the machine.
   4 """
   5
   6 from szyfrow.support.utilities import *
   7 from szyfrow.support.language_models import *
   8
   9
  10 class PocketEnigma(object):
  11     """A pocket enigma machine
  12     The wheel is internally represented as a 26-element list self.wheel_map,
  13     where wheel_map[i] == j shows that the position i places on from the arrow
  14     maps to the position j places on.
  15     """
  16     def __init__(self, wheel=1, position='a'):
  17         """initialise the pocket enigma, including which wheel to use and the
  18         starting position of the wheel.
  19
  20         The wheel is either 1 or 2 (the predefined wheels) or a list of letter
  21         pairs.
  22
  23         The position is the letter pointed to by the arrow on the wheel.
  24
  25         >>> pe.wheel_map
  26         [25, 4, 23, 10, 1, 7, 9, 5, 12, 6, 3, 17, 8, 14, 13, 21, 19, 11, 20, 16, 18, 15, 24, 2, 22, 0]
  27         >>> pe.position
  28         0
  29         """
  30         self.wheel1 = [('a', 'z'), ('b', 'e'), ('c', 'x'), ('d', 'k'),
  31             ('f', 'h'), ('g', 'j'), ('i', 'm'), ('l', 'r'), ('n', 'o'),
  32             ('p', 'v'), ('q', 't'), ('s', 'u'), ('w', 'y')]
  33         self.wheel2 = [('a', 'c'), ('b', 'd'), ('e', 'w'), ('f', 'i'),
  34             ('g', 'p'), ('h', 'm'), ('j', 'k'), ('l', 'n'), ('o', 'q'),
  35             ('r', 'z'), ('s', 'u'), ('t', 'v'), ('x', 'y')]
  36         if wheel == 1:
  37             self.make_wheel_map(self.wheel1)
  38         elif wheel == 2:
  39             self.make_wheel_map(self.wheel2)
  40         else:
  41             self.validate_wheel_spec(wheel)
  42             self.make_wheel_map(wheel)
  43         if position in string.ascii_lowercase:
  44             self.position = pos(position)
  45         else:
  46             self.position = position
  47
  48     def make_wheel_map(self, wheel_spec):
  49         """Expands a wheel specification from a list of letter-letter pairs
  50         into a full wheel_map.
  51
  52         >>> pe.make_wheel_map(pe.wheel2)
  53         [2, 3, 0, 1, 22, 8, 15, 12, 5, 10, 9, 13, 7, 11, 16, 6, 14, 25, 20, 21, 18, 19, 4, 24, 23, 17]
  54         """
  55         self.validate_wheel_spec(wheel_spec)
  56         self.wheel_map = [0] * 26
  57         for p in wheel_spec:
  58             self.wheel_map[pos(p[0])] = pos(p[1])
  59             self.wheel_map[pos(p[1])] = pos(p[0])
  60         return self.wheel_map
  61
  62     def validate_wheel_spec(self, wheel_spec):
  63         """Validates that a wheel specificaiton will turn into a valid wheel
  64         map.
  65
  66         >>> pe.validate_wheel_spec([])
  67         Traceback (most recent call last):
  68             ...
  69         ValueError: Wheel specification has 0 pairs, requires 13
  70         >>> pe.validate_wheel_spec([('a', 'b', 'c')]*13)
  71         Traceback (most recent call last):
  72             ...
  73         ValueError: Not all mappings in wheel specificationhave two elements
  74         >>> pe.validate_wheel_spec([('a', 'b')]*13)
  75         Traceback (most recent call last):
  76             ...
  77         ValueError: Wheel specification does not contain 26 letters
  78         """
  79         if len(wheel_spec) != 13:
  80             raise ValueError("Wheel specification has {} pairs, requires 13".
  81                 format(len(wheel_spec)))
  82         for p in wheel_spec:
  83             if len(p) != 2:
  84                 raise ValueError("Not all mappings in wheel specification"
  85                     "have two elements")
  86         if len(set([p[0] for p in wheel_spec] +
  87                     [p[1] for p in wheel_spec])) != 26:
  88             raise ValueError("Wheel specification does not contain 26 letters")
  89
  90     def encipher_letter(self, letter):
  91         """Enciphers a single letter, by advancing the wheel before looking up
  92         the letter on the wheel.
  93
  94         >>> pe.set_position('f')
  95         5
  96         >>> pe.encipher_letter('k')
  97         'h'
  98         """
  99         self.advance()
 100         return self.lookup(letter)
 101     decipher_letter = encipher_letter
 102
 103     def lookup(self, letter):
 104         """Look up what a letter enciphers to, without turning the wheel.
 105
 106         >>> pe.set_position('f')
 107         5
 108         >>> cat([pe.lookup(l) for l in string.ascii_lowercase])
 109         'udhbfejcpgmokrliwntsayqzvx'
 110         >>> pe.lookup('A')
 111         ''
 112         """
 113         if letter in string.ascii_lowercase:
 114             return unpos(
 115                 (self.wheel_map[(pos(letter) - self.position) % 26] +
 116                     self.position))
 117         else:
 118             return ''
 119
 120     def advance(self):
 121         """Advances the wheel one position.
 122
 123         >>> pe.set_position('f')
 124         5
 125         >>> pe.advance()
 126         6
 127         """
 128         self.position = (self.position + 1) % 26
 129         return self.position
 130
 131     def encipher(self, message, starting_position=None):
 132         """Enciphers a whole message.
 133
 134         >>> pe.set_position('f')
 135         5
 136         >>> pe.encipher('helloworld')
 137         'kjsglcjoqc'
 138         >>> pe.set_position('f')
 139         5
 140         >>> pe.encipher('kjsglcjoqc')
 141         'helloworld'
 142         >>> pe.encipher('helloworld', starting_position = 'x')
 143         'egrekthnnf'
 144         """
 145         if starting_position:
 146             self.set_position(starting_position)
 147         transformed = ''
 148         for l in message:
 149             transformed += self.encipher_letter(l)
 150         return transformed
 151     decipher = encipher
 152
 153     def set_position(self, position):
 154         """Sets the position of the wheel, by specifying the letter the arrow
 155         points to.
 156
 157         >>> pe.set_position('a')
 158         0
 159         >>> pe.set_position('m')
 160         12
 161         >>> pe.set_position('z')
 162         25
 163         """
 164         self.position = pos(position)
 165         return self.position
 166
 167
 168 def pocket_enigma_break_by_crib(message, wheel_spec, crib, crib_position):
 169     """Break a pocket enigma using a crib (some plaintext that's expected to
 170     be in a certain position). Returns a list of possible starting wheel
 171     positions that could produce the crib.
 172
 173     >>> pocket_enigma_break_by_crib('kzpjlzmoga', 1, 'h', 0)
 174     ['a', 'f', 'q']
 175     >>> pocket_enigma_break_by_crib('kzpjlzmoga', 1, 'he', 0)
 176     ['a']
 177     >>> pocket_enigma_break_by_crib('kzpjlzmoga', 1, 'll', 2)
 178     ['a']
 179     >>> pocket_enigma_break_by_crib('kzpjlzmoga', 1, 'l', 2)
 180     ['a']
 181     >>> pocket_enigma_break_by_crib('kzpjlzmoga', 1, 'l', 3)
 182     ['a', 'j', 'n']
 183     >>> pocket_enigma_break_by_crib('aaaaa', 1, 'l', 3)
 184     []
 185     """
 186     pe = PocketEnigma(wheel=wheel_spec)
 187     possible_positions = []
 188     for p in string.ascii_lowercase:
 189         pe.set_position(p)
 190         plaintext = pe.decipher(message)
 191         if plaintext[crib_position:crib_position+len(crib)] == crib:
 192             possible_positions += [p]
 193     return possible_positions
 194
 195 if __name__ == "__main__":
 196     import doctest
 197     doctest.testmod(extraglobs={'pe': PocketEnigma(1, 'a')})