szyfrow/enigma.py

   1 """A simulator for Enigma machines.
   2
   3 See `szyfrow.bombe.Bombe` for an implementation of the Bombe to break Enigma
   4 messages.
   5
   6 Specification from [Codes and Ciphers](http://www.codesandciphers.org.uk/enigma/rotorspec.htm) page.
   7
   8 Example Enigma machines from [Louise Dale](http://enigma.louisedade.co.uk/enigma.html) (full simulation) and [EnigmaCo](http://enigmaco.de/enigma/enigma.html) (good animation of the wheels, but no ring settings).
   9
  10 There's also the nice Enigma simulator for Android by [Franklin Heath](https://franklinheath.co.uk/2012/02/04/our-first-app-published-enigma-simulator/), available on the [Google Play store](https://play.google.com/store/apps/details?id=uk.co.franklinheath.enigmasim&hl=en_GB).
  11 """
  12
  13 import string
  14 import collections
  15 import multiprocessing
  16 import itertools
  17
  18 from szyfrow.support.utilities import *
  19
  20 # # Some convenience functions
  21
  22 # cat = ''.join
  23
  24 # def clean(text): return cat(l.lower() for l in text if l in string.ascii_letters)
  25
  26 # def pos(letter):
  27 #     if letter in string.ascii_lowercase:
  28 #         return ord(letter) - ord('a')
  29 #     elif letter in string.ascii_uppercase:
  30 #         return ord(letter) - ord('A')
  31 #     else:
  32 #         return ''
  33
  34 # def unpos(number): return chr(number % 26 + ord('a'))
  35
  36
  37 wheel_i_spec = 'ekmflgdqvzntowyhxuspaibrcj'
  38 wheel_ii_spec = 'ajdksiruxblhwtmcqgznpyfvoe'
  39 wheel_iii_spec = 'bdfhjlcprtxvznyeiwgakmusqo'
  40 wheel_iv_spec = 'esovpzjayquirhxlnftgkdcmwb'
  41 wheel_v_spec = 'vzbrgityupsdnhlxawmjqofeck'
  42 wheel_vi_spec = 'jpgvoumfyqbenhzrdkasxlictw'
  43 wheel_vii_spec = 'nzjhgrcxmyswboufaivlpekqdt'
  44 wheel_viii_spec = 'fkqhtlxocbjspdzramewniuygv'
  45 beta_wheel_spec = 'leyjvcnixwpbqmdrtakzgfuhos'
  46 gamma_wheel_spec = 'fsokanuerhmbtiycwlqpzxvgjd'
  47
  48 wheel_i_notches = ['q']
  49 wheel_ii_notches = ['e']
  50 wheel_iii_notches = ['v']
  51 wheel_iv_notches = ['j']
  52 wheel_v_notches = ['z']
  53 wheel_vi_notches = ['z', 'm']
  54 wheel_vii_notches = ['z', 'm']
  55 wheel_viii_notches = ['z', 'm']
  56
  57 reflector_b_spec = 'ay br cu dh eq fs gl ip jx kn mo tz vw'
  58 reflector_c_spec = 'af bv cp dj ei go hy kr lz mx nw tq su'
  59
  60
  61
  62 class LetterTransformer(object):
  63     """A generic substitution cipher, that has different transforms in the
  64     forward and backward directions. It requires that the transforms for all
  65     letters by provided.
  66
  67     A `transform` is a list of letter pairs, like `[('a', 'b'), ('c', 'd')]`.
  68     That would say that, in the forward direction `a` goes to `b` and
  69     `c` goes to `d`. In the backward direction, `b` goes to `a` and `d` goes
  70     to `c`.
  71     """
  72     def __init__(self, specification, raw_transform=False):
  73         """Validate and create a new transformer. The transform is parsed by
  74         `LetterTransformer.parse_specification` unless `raw_transform` is `True`
  75         """
  76         if raw_transform:
  77             transform = specification
  78         else:
  79             transform = self.parse_specification(specification)
  80         self.validate_transform(transform)
  81         self.make_transform_map(transform)
  82
  83     def parse_specification(self, specification):
  84         """Turn a `specification` string into a transform, by zipping it
  85         with ASCII lowercase letters to generate the pairs. This assumes that
  86         the `specification` defines the destination of the forward transform.
  87         """
  88         return list(zip(string.ascii_lowercase, sanitise(specification)))
  89         # return specification
  90
  91     def validate_transform(self, transform):
  92         """Checks that a transform is valid (every letter is mapped to
  93         exactly one other letter, in both directions).
  94         """
  95         if len(transform) != 26:
  96             raise ValueError("Transform specification has {} pairs, requires 26".
  97                 format(len(transform)))
  98         for p in transform:
  99             if len(p) != 2:
 100                 raise ValueError("Not all mappings in transform "
 101                     "have two elements")
 102         if len(set([p[0] for p in transform])) != 26:
 103             raise ValueError("Transform specification must list 26 origin letters")
 104         if len(set([p[1] for p in transform])) != 26:
 105             raise ValueError("Transform specification must list 26 destination letters")
 106
 107     def make_empty_transform(self):
 108         """An empty transform is one that maps every letter to 'a'.
 109         """
 110         self.forward_map = [0] * 26
 111         self.backward_map = [0] * 26
 112
 113     def make_transform_map(self, transform):
 114         """Create `forward_map` and `backward_map` from `transform`. The maps
 115         work on letter positions, not letter values. This makes the arithmetic
 116         for wheels much easier.
 117         """
 118         self.make_empty_transform()
 119         for p in transform:
 120             self.forward_map[pos(p[0])] = pos(p[1])
 121             self.backward_map[pos(p[1])] = pos(p[0])
 122         return self.forward_map, self.backward_map
 123
 124     def forward(self, letter):
 125         """Apply a map in the forward direction.
 126         """
 127         if letter in string.ascii_lowercase:
 128             return unpos(self.forward_map[pos(letter)])
 129         else:
 130             return ''
 131
 132     def backward(self, letter):
 133         """Apply a map in the backward direction.
 134         """
 135         if letter in string.ascii_lowercase:
 136             return unpos(self.backward_map[pos(letter)])
 137         else:
 138             return ''
 139
 140
 141 class Plugboard(LetterTransformer):
 142     """A plugboard, a type of letter transformer where forward and backward
 143     transforms are the same. If a letter isn't explicitly transformed, it is
 144     kept as it is.
 145     """
 146
 147     def parse_specification(self, specification):
 148         """Convert a specification into a transform. The specification is
 149         given as a list of letter pairs.
 150         """
 151         return [tuple(sanitise(p)) for p in specification.split()]
 152
 153     def validate_transform(self, transform):
 154         """A set of pairs, of from-to. Does not require all 26 letters
 155         are in the transform.
 156         """
 157         for p in transform:
 158             if len(p) != 2:
 159                 raise ValueError("Not all mappings in transform"
 160                     "have two elements")
 161
 162     def make_empty_transform(self):
 163         """An empty transform maps every letter to itself.
 164         """
 165         self.forward_map = list(range(26))
 166         self.backward_map = list(range(26))
 167
 168     def make_transform_map(self, transform):
 169         """Makes the maps for a plugboard. Ensures that if the pair ('a', 'b')
 170         is in the specification, the pair ('b', 'a') is also present.
 171         """
 172         expanded_transform = transform + [tuple(reversed(p)) for p in transform]
 173         return super(Plugboard, self).make_transform_map(expanded_transform)
 174
 175
 176
 177
 178 class Reflector(Plugboard):
 179     """A reflector is a plugboard that requires 13 transforms.
 180     The 'plugboard' superclass ensures that all 13 transforms are also applied
 181     in reverse, making 26 transforms in all.
 182     """
 183     def validate_transform(self, transform):
 184         if len(transform) != 13:
 185             raise ValueError("Reflector specification has {} pairs, requires 13".
 186                 format(len(transform)))
 187         if len(set([p[0] for p in transform] +
 188                     [p[1] for p in transform])) != 26:
 189             raise ValueError("Reflector specification does not contain 26 letters")
 190         try:
 191             super(Reflector, self).validate_transform(transform)
 192         except ValueError as v:
 193             raise ValueError("Not all mappings in reflector have two elements")
 194
 195
 196
 197
 198 class SimpleWheel(LetterTransformer):
 199     """A wheel is a transform that rotates.
 200
 201     Looking from the right, letters go in sequence a-b-c clockwise around the
 202     wheel.
 203
 204     The position of the wheel is the number of spaces anticlockwise the wheel
 205     has turned.
 206
 207     Letter inputs and outputs are given relative to the frame holding the wheel,
 208     so if the wheel is advanced three places, an input of 'p' will enter the
 209     wheel on the position under the wheel's 's' label.
 210     """
 211     def __init__(self, transform, position='a', raw_transform=False):
 212         super(SimpleWheel, self).__init__(transform, raw_transform)
 213         self.set_position(position)
 214
 215     def __getattribute__(self,name):
 216         if name=='position_l':
 217             return unpos(self.position)
 218         else:
 219             return object.__getattribute__(self, name)
 220
 221     def set_position(self, position):
 222         """Sets a wheel's position. If the `position` is a string, convert it
 223         to a number and set the position.
 224         """
 225         if isinstance(position, str):
 226             self.position = pos(position)
 227         else:
 228             self.position = position
 229
 230     def forward(self, letter):
 231         """Give the transformed letter in the forward direction, accounting
 232         for the position of the wheel.
 233         """
 234         if letter in string.ascii_lowercase:
 235             return unpos((self.forward_map[(pos(letter) + self.position) % 26] - self.position))
 236         else:
 237             return ''
 238
 239     def backward(self, letter):
 240         """Give the transformed letter in the backward direction, accounting
 241         for the position of the wheel.
 242         """
 243         if letter in string.ascii_lowercase:
 244             return unpos((self.backward_map[(pos(letter) + self.position) % 26] - self.position))
 245         else:
 246             return ''
 247
 248     def advance(self):
 249         """Advance a wheel one position."""
 250         self.position = (self.position + 1) % 26
 251
 252
 253
 254 class Wheel(SimpleWheel):
 255     """A wheel with a movable ring.
 256
 257     The ring holds the letters and the notches that turn other wheels. The core
 258     holds the wiring that does the transformation.
 259
 260     The ring position is how many steps the core is turned relative to the ring.
 261     This is one-based, so a ring setting of 1 means the core and ring are
 262     aligned.
 263
 264     The position of the wheel is the position of the core (the transforms)
 265     relative to the neutral position.
 266
 267     The position_l is the position of the ring, or what would be observed
 268     by the user of the Enigma machine.
 269
 270     The notch_positions are the number of advances of this wheel before it will
 271     advance the next wheel.
 272     """
 273     def __init__(self, transform, ring_notch_letters, ring_setting=1,
 274             position='a', raw_transform=False):
 275         self.ring_notch_letters = ring_notch_letters
 276         self.ring_setting = ring_setting
 277         super(Wheel, self).__init__(transform, position=position,
 278             raw_transform=raw_transform)
 279         self.set_position(position)
 280
 281     def __getattribute__(self,name):
 282         if name=='position_l':
 283             return unpos(self.position + self.ring_setting - 1)
 284         else:
 285             return object.__getattribute__(self, name)
 286
 287     def set_position(self, position):
 288         if isinstance(position, str):
 289             self.position = (pos(position) - self.ring_setting + 1) % 26
 290         else:
 291             self.position = (position - self.ring_setting) % 26
 292         # # self.notch_positions = [(pos(p) - pos(position)) % 26  for p in self.ring_notch_letters]
 293         # self.notch_positions = [(pos(p) - (self.position + self.ring_setting - 1)) % 26  for p in self.ring_notch_letters]
 294         self.notch_positions = [(self.position + self.ring_setting - 1 - pos(p)) % 26  for p in self.ring_notch_letters]
 295
 296     def advance(self):
 297         """Advance a wheel's core, then advance the ring position to match.
 298         """
 299         super(Wheel, self).advance()
 300         self.notch_positions = [(p + 1) % 26 for p in self.notch_positions]
 301         return self.position
 302
 303
 304 class Enigma(object):
 305     """An Enigma machine.
 306
 307
 308     """
 309     def __init__(self, reflector_spec,
 310                  left_wheel_spec, left_wheel_notches,
 311                  middle_wheel_spec, middle_wheel_notches,
 312                  right_wheel_spec, right_wheel_notches,
 313                  left_ring_setting, middle_ring_setting, right_ring_setting,
 314                  plugboard_setting):
 315         self.reflector = Reflector(reflector_spec)
 316         self.left_wheel = Wheel(left_wheel_spec, left_wheel_notches,
 317             ring_setting=left_ring_setting)
 318         self.middle_wheel = Wheel(middle_wheel_spec, middle_wheel_notches,
 319             ring_setting=middle_ring_setting)
 320         self.right_wheel = Wheel(right_wheel_spec, right_wheel_notches,
 321             ring_setting=right_ring_setting)
 322         self.plugboard = Plugboard(plugboard_setting)
 323
 324     def __getattribute__(self,name):
 325         if name=='wheel_positions':
 326             return (self.left_wheel.position,
 327                     self.middle_wheel.position,
 328                     self.right_wheel.position
 329                     )
 330         elif name=='wheel_positions_l':
 331             return (self.left_wheel.position_l,
 332                     self.middle_wheel.position_l,
 333                     self.right_wheel.position_l
 334                     )
 335         elif name=='notch_positions':
 336             return (self.left_wheel.notch_positions,
 337                     self.middle_wheel.notch_positions,
 338                     self.right_wheel.notch_positions
 339                     )
 340         else:
 341             return object.__getattribute__(self, name)
 342
 343     def set_wheels(self, left_wheel_position, middle_wheel_position,
 344             right_wheel_position):
 345         """Set the Enigma's wheels to the specified positions.
 346         """
 347         self.left_wheel.set_position(left_wheel_position)
 348         self.middle_wheel.set_position(middle_wheel_position)
 349         self.right_wheel.set_position(right_wheel_position)
 350
 351     def lookup(self, letter):
 352         """Lookup the enciphering of a letter, without advancing any wheels
 353         """
 354         a = self.plugboard.forward(letter)
 355         b = self.right_wheel.forward(a)
 356         c = self.middle_wheel.forward(b)
 357         d = self.left_wheel.forward(c)
 358         e = self.reflector.forward(d)
 359         f = self.left_wheel.backward(e)
 360         g = self.middle_wheel.backward(f)
 361         h = self.right_wheel.backward(g)
 362         i = self.plugboard.backward(h)
 363         return i
 364
 365     def advance(self):
 366         """Advance the Enigma's wheels one step. The right wheel always
 367         advances. The middle and right wheels may advance if the notches
 368         line up correctly.
 369         """
 370         advance_middle = False
 371         advance_left = False
 372         if 0 in self.right_wheel.notch_positions:
 373             advance_middle = True
 374         if 0 in self.middle_wheel.notch_positions:
 375             advance_left = True
 376             advance_middle = True
 377         self.right_wheel.advance()
 378         if advance_middle: self.middle_wheel.advance()
 379         if advance_left: self.left_wheel.advance()
 380
 381     def encipher_letter(self, letter):
 382         """Encipher a letter. Advance the Enigma machine, then lookup the
 383         encryption of a letter.
 384         """
 385         self.advance()
 386         return self.lookup(letter)
 387
 388     def encipher(self, message):
 389         """Encipher a message."""
 390         enciphered = ''
 391         for letter in sanitise(message):
 392             enciphered += self.encipher_letter(letter)
 393         return enciphered
 394
 395     decipher = encipher
 396
 397
 398 # for i in range(26):
 399 #     enigma.advance()
 400 #     print('enigma.advance()')
 401 #     print("assert(enigma.wheel_positions == {})".format(enigma.wheel_positions))
 402 #     print("assert(cat(enigma.wheel_positions_l) == '{}')".format(cat(enigma.wheel_positions_l)))
 403 #     print("assert(enigma.notch_positions == {})".format(enigma.notch_positions))
 404 #     print("assert(cat(enigma.lookup(l) for l in string.ascii_lowercase) == '{}')".format(cat(enigma.lookup(l) for l in string.ascii_lowercase)))
 405 #     print()
 406
 407
 408 if __name__ == "__main__":
 409     import doctest
 410     # doctest.testmod(extraglobs={'lt': LetterTransformer(1, 'a')})
 411     doctest.testmod()
 412