changeset 12:778f0bd51818

Add a method to retrieve settings from a M209 instance.
author Brian Neal <bgneal@gmail.com>
date Wed, 05 Jun 2013 20:41:54 -0500
parents 45af97a43ee1
children e242ef3b5a84
files m209/converter.py m209/key_wheel.py m209/tests/test_converter.py
diffstat 3 files changed, 33 insertions(+), 25 deletions(-) [+]
line wrap: on
line diff
--- a/m209/converter.py	Tue Jun 04 19:31:37 2013 -0500
+++ b/m209/converter.py	Wed Jun 05 20:41:54 2013 -0500
@@ -6,6 +6,7 @@
 assemble a complete M-209 converter.
 
 """
+from collections import namedtuple
 import string
 
 from . import M209Error
@@ -17,19 +18,19 @@
 CIPHER_TABLE = list(reversed(string.ascii_uppercase))
 
 
+M209Settings = namedtuple('M209Settings', ['lugs', 'pin_list'])
+
 class M209:
     """The M209 class is the top-level class in the M-209 simulation. It
     aggregates key wheels and a drum and orchestrates their movements to provide
     encrypt and decrypt functions for the operator.
 
     """
-    def __init__(self):
-        """Build a M209 instance with all pins in the ineffective state and all
-        drum lugs in neutral positions.
-
-        """
+    def __init__(self, lugs=None, pin_list=None):
+        """Build a M209 instance with the given lug & pin settings."""
         self.key_wheels = [KeyWheel(*args) for args in KEY_WHEEL_DATA]
-        self.drum = Drum()
+        self.set_drum_lugs(lugs)
+        self.set_all_pins(pin_list)
         self.letter_counter = 0
 
     def set_pins(self, n, effective_pins):
@@ -118,6 +119,12 @@
         for n in range(6):
             self.key_wheels[n].set_pos(s[n])
 
+    def get_settings(self):
+        """Returns the current settings as a M209Settings named tuple."""
+
+        return M209Settings(lugs=self.drum.key_list,
+                        pin_list=[kw.effective_pins for kw in self.key_wheels])
+
     def encrypt(self, plaintext, group=True, spaces=True):
         """Performs an encrypt operation on the given plaintext and returns
         the ciphertext as a string.
--- a/m209/key_wheel.py	Tue Jun 04 19:31:37 2013 -0500
+++ b/m209/key_wheel.py	Wed Jun 05 20:41:54 2013 -0500
@@ -62,6 +62,7 @@
     def reset_pins(self):
         """Reset all pins to the ineffective state."""
         self.pins = [False] * self.num_pins
+        self.effective_pins = ''
 
     def set_pins(self, effective_pins):
         """Sets which pins are effective.
@@ -73,6 +74,9 @@
 
         """
         self.reset_pins()
+        if not effective_pins:
+            return
+
         for letter in effective_pins:
             try:
                 n = self.letter_offsets[letter]
@@ -80,6 +84,8 @@
                 raise KeyWheelError("Invalid pin: {}".format(letter))
             self.pins[n] = True
 
+        self.effective_pins = effective_pins
+
     def rotate(self, steps=1):
         """Rotate the key wheel the given number of steps."""
         self.pos = (self.pos + steps) % self.num_pins
--- a/m209/tests/test_converter.py	Tue Jun 04 19:31:37 2013 -0500
+++ b/m209/tests/test_converter.py	Wed Jun 05 20:41:54 2013 -0500
@@ -11,7 +11,7 @@
 
 
 # Data taken from Mark J. Blair's AA key list
-AA_LUGS = '1-0*5 0-3*3 0-4 0-5*4 0-6*6 1-2 1-5*4 3-4 3-6 5-6'
+AA_LUGS = '0-4 0-5*4 0-6*6 1-0*5 1-2 1-5*4 3-0*3 3-4 3-6 5-6'
 
 AA_PIN_LIST = [
     'FGIKOPRSUVWYZ',
@@ -52,13 +52,9 @@
         pt = 'A' * 26
         ct = check
 
-        m = M209()
-        m.set_drum_lugs(lugs)
-
-        m.set_all_pins(pin_list)
+        m = M209(lugs, pin_list)
 
         result = m.encrypt(pt)
-
         self.assertEqual(result, ct)
         self.assertEqual(m.letter_counter, 26)
 
@@ -93,10 +89,7 @@
 
     def test_no_group(self):
 
-        m = M209()
-        m.set_drum_lugs(AA_LUGS)
-        m.set_all_pins(AA_PIN_LIST)
-
+        m = M209(AA_LUGS, AA_PIN_LIST)
         result = m.encrypt('A' * 26, group=False)
         expected = AA_CHECK.replace(' ', '')
         self.assertEqual(result, expected)
@@ -108,9 +101,7 @@
 
     def test_encrypt_spaces(self):
 
-        m = M209()
-        m.set_drum_lugs(AA_LUGS)
-        m.set_all_pins(AA_PIN_LIST)
+        m = M209(AA_LUGS, AA_PIN_LIST)
 
         wheels = 'YGXREL'
         m.set_key_wheels(wheels)
@@ -131,9 +122,7 @@
 
     def test_decrypt_no_z_sub(self):
 
-        m = M209()
-        m.set_drum_lugs(AA_LUGS)
-        m.set_all_pins(AA_PIN_LIST)
+        m = M209(AA_LUGS, AA_PIN_LIST)
 
         pt = 'ATTACK AT DAWN'
         wheels = 'YGXREL'
@@ -147,9 +136,7 @@
 
     def test_set_pins_vs_all_pins(self):
 
-        m1 = M209()
-        m1.set_drum_lugs(AA_LUGS)
-        m1.set_all_pins(AA_PIN_LIST)
+        m1 = M209(AA_LUGS, AA_PIN_LIST)
 
         pt = 'ATTACK AT DAWN'
         wheels = 'YGXREL'
@@ -165,3 +152,11 @@
         ct2 = m2.encrypt(pt)
 
         self.assertEqual(ct1, ct2)
+
+    def test_get_settings(self):
+
+        m = M209(AA_LUGS, AA_PIN_LIST)
+        settings = m.get_settings()
+
+        self.assertEqual(settings.lugs, AA_LUGS)
+        self.assertEqual(settings.pin_list, AA_PIN_LIST)