Started on documentation
[szyfrow.git] / szyfrow / enigma.py
index 89758f532087985078041473006ce3d6510e3933..2caa397c53c82ad9768f8b950f0d273e7554cb7f 100644 (file)
@@ -1,52 +1,37 @@
+"""A simulator for Enigma machines.
 
 
-# coding: utf-8
+See `szyfrow.bombe.Bombe` for an implementation of the Bombe to break Enigma
+messages.
 
 
-##################################
-# # Enigma machine
-##################################
-# Specification from [Codes and Ciphers](http://www.codesandciphers.org.uk/enigma/rotorspec.htm) page.
-# 
-# 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).
-# 
-# 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).
+Specification from [Codes and Ciphers](http://www.codesandciphers.org.uk/enigma/rotorspec.htm) page.
 
 
+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).
 
 
+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).
+"""
 
 import string
 import collections
 import multiprocessing
 import itertools
 
 import string
 import collections
 import multiprocessing
 import itertools
-import logging
 
 
-logger = logging.getLogger('engima')
-logger.setLevel(logging.WARNING)
-# logger.setLevel(logging.INFO)
-# logger.setLevel(logging.DEBUG)
+from szyfrow.support.utilities import *
 
 
-# create the logging file handler
-fh = logging.FileHandler("enigma.log")
-formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-fh.setFormatter(formatter)
+# # Some convenience functions
 
 
-# add handler to logger object
-logger.addHandler(fh)
+# cat = ''.join
 
 
+# def clean(text): return cat(l.lower() for l in text if l in string.ascii_letters)
 
 
-# Some convenience functions
-
-cat = ''.join
-
-def clean(text): return cat(l.lower() for l in text if l in string.ascii_letters)
-
-def pos(letter): 
-    if letter in string.ascii_lowercase:
-        return ord(letter) - ord('a')
-    elif letter in string.ascii_uppercase:
-        return ord(letter) - ord('A')
-    else:
-        return ''
+# def pos(letter): 
+#     if letter in string.ascii_lowercase:
+#         return ord(letter) - ord('a')
+#     elif letter in string.ascii_uppercase:
+#         return ord(letter) - ord('A')
+#     else:
+#         return ''
     
     
-def unpos(number): return chr(number % 26 + ord('a'))
+def unpos(number): return chr(number % 26 + ord('a'))
 
 
 wheel_i_spec = 'ekmflgdqvzntowyhxuspaibrcj'
 
 
 wheel_i_spec = 'ekmflgdqvzntowyhxuspaibrcj'
@@ -78,8 +63,16 @@ class LetterTransformer(object):
     """A generic substitution cipher, that has different transforms in the 
     forward and backward directions. It requires that the transforms for all
     letters by provided.
     """A generic substitution cipher, that has different transforms in the 
     forward and backward directions. It requires that the transforms for all
     letters by provided.
+
+    A `transform` is a list of letter pairs, like `[('a', 'b'), ('c', 'd')]`.
+    That would say that, in the forward direction `a` goes to `b` and 
+    `c` goes to `d`. In the backward direction, `b` goes to `a` and `d` goes
+    to `c`. 
     """
     def __init__(self, specification, raw_transform=False):
     """
     def __init__(self, specification, raw_transform=False):
+        """Validate and create a new transformer. The transform is parsed by
+        `LetterTransformer.parse_specification` unless `raw_transform` is `True`
+        """
         if raw_transform:
             transform = specification
         else:
         if raw_transform:
             transform = specification
         else:
@@ -88,11 +81,17 @@ class LetterTransformer(object):
         self.make_transform_map(transform)
     
     def parse_specification(self, specification):
         self.make_transform_map(transform)
     
     def parse_specification(self, specification):
-        return list(zip(string.ascii_lowercase, clean(specification)))
+        """Turn a `specification` string into a transform, by zipping it
+        with ASCII lowercase letters to generate the pairs. This assumes that
+        the `specification` defines the destination of the forward transform.
+        """
+        return list(zip(string.ascii_lowercase, sanitise(specification)))
         # return specification
     
     def validate_transform(self, transform):
         # return specification
     
     def validate_transform(self, transform):
-        """A set of pairs, of from-to"""
+        """Checks that a transform is valid (every letter is mapped to 
+        exactly one other letter, in both directions).
+        """
         if len(transform) != 26:
             raise ValueError("Transform specification has {} pairs, requires 26".
                 format(len(transform)))
         if len(transform) != 26:
             raise ValueError("Transform specification has {} pairs, requires 26".
                 format(len(transform)))
@@ -106,10 +105,16 @@ class LetterTransformer(object):
             raise ValueError("Transform specification must list 26 destination letters") 
 
     def make_empty_transform(self):
             raise ValueError("Transform specification must list 26 destination letters") 
 
     def make_empty_transform(self):
+        """An empty transform is one that maps every letter to 'a'.
+        """
         self.forward_map = [0] * 26
         self.backward_map = [0] * 26
             
     def make_transform_map(self, transform):
         self.forward_map = [0] * 26
         self.backward_map = [0] * 26
             
     def make_transform_map(self, transform):
+        """Create `forward_map` and `backward_map` from `transform`. The maps
+        work on letter positions, not letter values. This makes the arithmetic
+        for wheels much easier.
+        """
         self.make_empty_transform()
         for p in transform:
             self.forward_map[pos(p[0])] = pos(p[1])
         self.make_empty_transform()
         for p in transform:
             self.forward_map[pos(p[0])] = pos(p[1])
@@ -117,12 +122,16 @@ class LetterTransformer(object):
         return self.forward_map, self.backward_map
     
     def forward(self, letter):
         return self.forward_map, self.backward_map
     
     def forward(self, letter):
+        """Apply a map in the forward direction.
+        """
         if letter in string.ascii_lowercase:
             return unpos(self.forward_map[pos(letter)])
         else:
             return ''
                 
     def backward(self, letter):
         if letter in string.ascii_lowercase:
             return unpos(self.forward_map[pos(letter)])
         else:
             return ''
                 
     def backward(self, letter):
+        """Apply a map in the backward direction.
+        """
         if letter in string.ascii_lowercase:
             return unpos(self.backward_map[pos(letter)])
         else:
         if letter in string.ascii_lowercase:
             return unpos(self.backward_map[pos(letter)])
         else:
@@ -134,21 +143,32 @@ class Plugboard(LetterTransformer):
     transforms are the same. If a letter isn't explicitly transformed, it is 
     kept as it is.
     """
     transforms are the same. If a letter isn't explicitly transformed, it is 
     kept as it is.
     """
+
     def parse_specification(self, specification):
     def parse_specification(self, specification):
-        return [tuple(clean(p)) for p in specification.split()]
+        """Convert a specification into a transform. The specification is
+        given as a list of letter pairs.
+        """
+        return [tuple(sanitise(p)) for p in specification.split()]
     
     def validate_transform(self, transform):
     
     def validate_transform(self, transform):
-        """A set of pairs, of from-to"""
+        """A set of pairs, of from-to. Does not require all 26 letters
+        are in the transform.
+        """
         for p in transform:
             if len(p) != 2:
                 raise ValueError("Not all mappings in transform"
                     "have two elements")
     
     def make_empty_transform(self):
         for p in transform:
             if len(p) != 2:
                 raise ValueError("Not all mappings in transform"
                     "have two elements")
     
     def make_empty_transform(self):
+        """An empty transform maps every letter to itself.
+        """
         self.forward_map = list(range(26))
         self.backward_map = list(range(26))
         
     def make_transform_map(self, transform):
         self.forward_map = list(range(26))
         self.backward_map = list(range(26))
         
     def make_transform_map(self, transform):
+        """Makes the maps for a plugboard. Ensures that if the pair ('a', 'b')
+        is in the specification, the pair ('b', 'a') is also present.
+        """
         expanded_transform = transform + [tuple(reversed(p)) for p in transform]
         return super(Plugboard, self).make_transform_map(expanded_transform)
 
         expanded_transform = transform + [tuple(reversed(p)) for p in transform]
         return super(Plugboard, self).make_transform_map(expanded_transform)
 
@@ -157,6 +177,8 @@ class Plugboard(LetterTransformer):
 
 class Reflector(Plugboard):
     """A reflector is a plugboard that requires 13 transforms.
 
 class Reflector(Plugboard):
     """A reflector is a plugboard that requires 13 transforms.
+    The 'plugboard' superclass ensures that all 13 transforms are also applied
+    in reverse, making 26 transforms in all.
     """
     def validate_transform(self, transform):
         if len(transform) != 13:
     """
     def validate_transform(self, transform):
         if len(transform) != 13:
@@ -184,7 +206,7 @@ class SimpleWheel(LetterTransformer):
 
     Letter inputs and outputs are given relative to the frame holding the wheel,
     so if the wheel is advanced three places, an input of 'p' will enter the 
 
     Letter inputs and outputs are given relative to the frame holding the wheel,
     so if the wheel is advanced three places, an input of 'p' will enter the 
-    wheel on the position under the wheel's 'q' label.
+    wheel on the position under the wheel's 's' label.
     """
     def __init__(self, transform, position='a', raw_transform=False):
         super(SimpleWheel, self).__init__(transform, raw_transform)
     """
     def __init__(self, transform, position='a', raw_transform=False):
         super(SimpleWheel, self).__init__(transform, raw_transform)
@@ -197,25 +219,34 @@ class SimpleWheel(LetterTransformer):
             return object.__getattribute__(self, name)
     
     def set_position(self, position):
             return object.__getattribute__(self, name)
     
     def set_position(self, position):
+        """Sets a wheel's position. If the `position` is a string, convert it
+        to a number and set the position.
+        """
         if isinstance(position, str):
         if isinstance(position, str):
-            # self.position = ord(position) - ord('a')
             self.position = pos(position)
         else:
             self.position = position
     
     def forward(self, letter):
             self.position = pos(position)
         else:
             self.position = position
     
     def forward(self, letter):
+        """Give the transformed letter in the forward direction, accounting
+        for the position of the wheel.
+        """
         if letter in string.ascii_lowercase:
             return unpos((self.forward_map[(pos(letter) + self.position) % 26] - self.position))
         else:
             return ''
                 
     def backward(self, letter):
         if letter in string.ascii_lowercase:
             return unpos((self.forward_map[(pos(letter) + self.position) % 26] - self.position))
         else:
             return ''
                 
     def backward(self, letter):
+        """Give the transformed letter in the backward direction, accounting
+        for the position of the wheel.
+        """
         if letter in string.ascii_lowercase:
             return unpos((self.backward_map[(pos(letter) + self.position) % 26] - self.position))
         else:
             return ''
         
     def advance(self):
         if letter in string.ascii_lowercase:
             return unpos((self.backward_map[(pos(letter) + self.position) % 26] - self.position))
         else:
             return ''
         
     def advance(self):
+        """Advance a wheel one position."""
         self.position = (self.position + 1) % 26
 
 
         self.position = (self.position + 1) % 26
 
 
@@ -238,12 +269,13 @@ class Wheel(SimpleWheel):
 
     The notch_positions are the number of advances of this wheel before it will 
     advance the next wheel.
 
     The notch_positions are the number of advances of this wheel before it will 
     advance the next wheel.
-
     """
     """
-    def __init__(self, transform, ring_notch_letters, ring_setting=1, position='a', raw_transform=False):
+    def __init__(self, transform, ring_notch_letters, ring_setting=1, 
+            position='a', raw_transform=False):
         self.ring_notch_letters = ring_notch_letters
         self.ring_setting = ring_setting
         self.ring_notch_letters = ring_notch_letters
         self.ring_setting = ring_setting
-        super(Wheel, self).__init__(transform, position=position, raw_transform=raw_transform)
+        super(Wheel, self).__init__(transform, position=position, 
+            raw_transform=raw_transform)
         self.set_position(position)
         
     def __getattribute__(self,name):
         self.set_position(position)
         
     def __getattribute__(self,name):
@@ -262,6 +294,8 @@ class Wheel(SimpleWheel):
         self.notch_positions = [(self.position + self.ring_setting - 1 - pos(p)) % 26  for p in self.ring_notch_letters]
         
     def advance(self):
         self.notch_positions = [(self.position + self.ring_setting - 1 - pos(p)) % 26  for p in self.ring_notch_letters]
         
     def advance(self):
+        """Advance a wheel's core, then advance the ring position to match.
+        """
         super(Wheel, self).advance()
         self.notch_positions = [(p + 1) % 26 for p in self.notch_positions]
         return self.position
         super(Wheel, self).advance()
         self.notch_positions = [(p + 1) % 26 for p in self.notch_positions]
         return self.position
@@ -279,27 +313,44 @@ class Enigma(object):
                  left_ring_setting, middle_ring_setting, right_ring_setting,
                  plugboard_setting):
         self.reflector = Reflector(reflector_spec)
                  left_ring_setting, middle_ring_setting, right_ring_setting,
                  plugboard_setting):
         self.reflector = Reflector(reflector_spec)
-        self.left_wheel = Wheel(left_wheel_spec, left_wheel_notches, ring_setting=left_ring_setting)
-        self.middle_wheel = Wheel(middle_wheel_spec, middle_wheel_notches, ring_setting=middle_ring_setting)
-        self.right_wheel = Wheel(right_wheel_spec, right_wheel_notches, ring_setting=right_ring_setting)
+        self.left_wheel = Wheel(left_wheel_spec, left_wheel_notches, 
+            ring_setting=left_ring_setting)
+        self.middle_wheel = Wheel(middle_wheel_spec, middle_wheel_notches, 
+            ring_setting=middle_ring_setting)
+        self.right_wheel = Wheel(right_wheel_spec, right_wheel_notches, 
+            ring_setting=right_ring_setting)
         self.plugboard = Plugboard(plugboard_setting)
         
     def __getattribute__(self,name):
         if name=='wheel_positions':
         self.plugboard = Plugboard(plugboard_setting)
         
     def __getattribute__(self,name):
         if name=='wheel_positions':
-            return self.left_wheel.position, self.middle_wheel.position, self.right_wheel.position 
+            return (self.left_wheel.position, 
+                    self.middle_wheel.position, 
+                    self.right_wheel.position
+                    )
         elif name=='wheel_positions_l':
         elif name=='wheel_positions_l':
-            return self.left_wheel.position_l, self.middle_wheel.position_l, self.right_wheel.position_l 
+            return (self.left_wheel.position_l, 
+                    self.middle_wheel.position_l, 
+                    self.right_wheel.position_l
+                    )
         elif name=='notch_positions':
         elif name=='notch_positions':
-            return self.left_wheel.notch_positions, self.middle_wheel.notch_positions, self.right_wheel.notch_positions
+            return (self.left_wheel.notch_positions, 
+                    self.middle_wheel.notch_positions, 
+                    self.right_wheel.notch_positions
+                    )
         else:
             return object.__getattribute__(self, name)
 
         else:
             return object.__getattribute__(self, name)
 
-    def set_wheels(self, left_wheel_position, middle_wheel_position, right_wheel_position):
+    def set_wheels(self, left_wheel_position, middle_wheel_position, 
+            right_wheel_position):
+        """Set the Enigma's wheels to the specified positions.
+        """
         self.left_wheel.set_position(left_wheel_position)
         self.middle_wheel.set_position(middle_wheel_position)
         self.right_wheel.set_position(right_wheel_position)
         
     def lookup(self, letter):
         self.left_wheel.set_position(left_wheel_position)
         self.middle_wheel.set_position(middle_wheel_position)
         self.right_wheel.set_position(right_wheel_position)
         
     def lookup(self, letter):
+        """Lookup the enciphering of a letter, without advancing any wheels
+        """
         a = self.plugboard.forward(letter)
         b = self.right_wheel.forward(a)
         c = self.middle_wheel.forward(b)
         a = self.plugboard.forward(letter)
         b = self.right_wheel.forward(a)
         c = self.middle_wheel.forward(b)
@@ -312,6 +363,10 @@ class Enigma(object):
         return i
     
     def advance(self):
         return i
     
     def advance(self):
+        """Advance the Enigma's wheels one step. The right wheel always
+        advances. The middle and right wheels may advance if the notches
+        line up correctly.
+        """
         advance_middle = False
         advance_left = False
         if 0 in self.right_wheel.notch_positions:
         advance_middle = False
         advance_left = False
         if 0 in self.right_wheel.notch_positions:
@@ -324,12 +379,16 @@ class Enigma(object):
         if advance_left: self.left_wheel.advance()
             
     def encipher_letter(self, letter):
         if advance_left: self.left_wheel.advance()
             
     def encipher_letter(self, letter):
+        """Encipher a letter. Advance the Enigma machine, then lookup the
+        encryption of a letter.
+        """
         self.advance()
         return self.lookup(letter)
     
     def encipher(self, message):
         self.advance()
         return self.lookup(letter)
     
     def encipher(self, message):
+        """Encipher a message."""
         enciphered = ''
         enciphered = ''
-        for letter in clean(message):
+        for letter in sanitise(message):
             enciphered += self.encipher_letter(letter)
         return enciphered
 
             enciphered += self.encipher_letter(letter)
         return enciphered