changeset 14:0fe1c4a11bad

Created a factory function to make it easier to create Enigmas.
author Brian Neal <bgneal@gmail.com>
date Sun, 27 May 2012 15:05:00 -0500
parents 3fbdc7005075
children 7ea7da689dbd
files enigma/machine.py enigma/main.py enigma/tests/test_enigma.py
diffstat 3 files changed, 44 insertions(+), 28 deletions(-) [+]
line wrap: on
line diff
--- a/enigma/machine.py	Sun May 27 13:57:13 2012 -0500
+++ b/enigma/machine.py	Sun May 27 15:05:00 2012 -0500
@@ -8,6 +8,10 @@
 """
 import string
 
+from .rotors.factory import create_rotor, create_reflector
+from .plugboard import Plugboard
+
+
 class EnigmaError(Exception):
     pass
 
@@ -46,6 +50,40 @@
         self.reflector = reflector
         self.plugboard = plugboard
 
+    @classmethod
+    def from_key_sheet(cls, rotors, ring_settings=(0, 0, 0), reflector='B',
+            plugboard_settings=''):
+        """Convenience function to build an EnigmaMachine from the data as you
+        might find it on a key sheet:
+
+        rotors: a list of strings naming the rotors from left to right;
+            e.g. ["I", "III", "IV"]
+
+        ring_settings: an iterable of integers representing the ring settings to
+        be applied to the rotors in the rotors list
+
+        reflector: a string that names the reflector to use
+
+        plugboard: a string of plugboard settings as you might find on a key
+        sheet; e.g. 'PO ML IU KJ NH YT GB VF RE DC' 
+
+        """
+        # validate inputs
+        num_rotors = len(rotors)
+        if num_rotors not in (3, 4):
+            raise EnigmaError("invalid rotors list size")
+
+        if num_rotors != len(ring_settings):
+            raise EnigmaError("please provide %d ring settings" % num_rotors)
+
+        # assemble the machine
+        rotor_list = [create_rotor(r[0], r[1]) for r in zip(rotors, ring_settings)]
+
+        return cls(rotor_list, 
+                   create_reflector(reflector),
+                   Plugboard(plugboard_settings))
+
+
     def set_display(self, val):
         """Sets the rotor operator windows to 'val'.
 
--- a/enigma/main.py	Sun May 27 13:57:13 2012 -0500
+++ b/enigma/main.py	Sun May 27 15:05:00 2012 -0500
@@ -2,21 +2,14 @@
 # This file is part of Py-Enigma, the Enigma Machine simulation.
 # Py-Enigma is released under the MIT License (see License.txt).
 
+from .machine import EnigmaMachine
 
-from rotors.factory import create_rotor
-from rotors.factory import create_reflector
-from machine import EnigmaMachine
 
 def main():
 
-    rotors = []
-    rotors.append(create_rotor('I'))
-    rotors.append(create_rotor('II'))
-    rotors.append(create_rotor('III'))
-
-    reflector = create_reflector('B')
-
-    machine = EnigmaMachine(rotors=rotors, reflector=reflector)
+    machine = EnigmaMachine.from_key_sheet(
+                        rotors=['I', 'II', 'III'],
+                        reflector='B')
 
     machine.set_display('AAA')
     cipher_text = machine.process_text('AAAAA')
--- a/enigma/tests/test_enigma.py	Sun May 27 13:57:13 2012 -0500
+++ b/enigma/tests/test_enigma.py	Sun May 27 15:05:00 2012 -0500
@@ -6,9 +6,7 @@
 
 import unittest
 
-from ..rotors.factory import create_rotor, create_reflector
 from ..machine import EnigmaMachine
-from ..plugboard import Plugboard
 
 
 class SteppingTestCase(unittest.TestCase):
@@ -21,12 +19,7 @@
         # This example taken from 
         # http://users.telenet.be/d.rijmenants/en/enigmatech.htm
         # in the section on "The Stepping Mechanism."
-        rotors = []
-        rotors.append(create_rotor("III"))
-        rotors.append(create_rotor("II"))
-        rotors.append(create_rotor("I"))
-
-        m = EnigmaMachine(rotors, create_reflector('B'), Plugboard())
+        m = EnigmaMachine.from_key_sheet(rotors=['III', 'II', 'I'])
 
         m.set_display('KDO')
 
@@ -43,15 +36,7 @@
     CIPHER_TEXT = 'BDZGO'
 
     def setUp(self):
-        rotors = []
-        rotors.append(create_rotor('I'))
-        rotors.append(create_rotor('II'))
-        rotors.append(create_rotor('III'))
-
-        reflector = create_reflector('B')
-
-        self.machine = EnigmaMachine(rotors=rotors, reflector=reflector,
-                                     plugboard=Plugboard())
+        self.machine = EnigmaMachine.from_key_sheet(rotors=['I', 'II', 'III'])
         self.machine.set_display('AAA')
 
     def test_simple_encrypt(self):