szyfrow/pocket_enigma.py

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