changeset 16:1448d698f9e4

Refactor the encrypt procedure into a class.
author Brian Neal <bgneal@gmail.com>
date Sat, 08 Jun 2013 12:55:07 -0500
parents 2e2692fb7de6
children c8027f88443f
files m209/procedure.py m209/tests/test_procedure.py
diffstat 2 files changed, 101 insertions(+), 93 deletions(-) [+]
line wrap: on
line diff
--- a/m209/procedure.py	Fri Jun 07 22:30:33 2013 -0500
+++ b/m209/procedure.py	Sat Jun 08 12:55:07 2013 -0500
@@ -21,7 +21,7 @@
 import random
 
 from . import M209Error
-from .converter import M209_ALPHABET_SET, M209_ALPHABET_LIST
+from .converter import M209, M209_ALPHABET_SET, M209_ALPHABET_LIST
 from .key_wheel import KeyWheelError
 
 
@@ -29,108 +29,120 @@
     pass
 
 
-def encrypt(m_209, plaintext, group=True, spaces=True, key_list=None,
-            key_list_ind=None, ext_msg_ind=None, sys_ind=None):
-    """Encrypts a plaintext message using standard procedure. The encrypted text
-    with the required message indicators are returned as one string.
+class StdEncryptProcedure:
+    """This class encapsulates the "standard" encrypt procedure for the M-209 as
+    found in the training film T.F. 11 - 1400.
 
-    The encrypt function accepts these parameters:
-
-    m_209 - A M209 converter instance
-    plaintext - Input string of text to be encrypted
-    group - If True, the resulting encrypted text will be grouped into 5-letter
-        groups with a space between each group. If False, no spaces will be
-        present in the output.
-    spaces - If True, space characters in the input plaintext will automatically
-        be replaced with 'Z' characters before encrypting.
-    key_list - If not None, this must be a KeyList instance and it will be
-        used to setup the m_209 machine. If None, it is assumed the M209 is
-        already setup in the desired configuration.
-    key_list_ind - If key_list is None, then this parameter must be supplied,
-        and it is the key list indicator that was used to setup the M209.
-    ext_msg_ind - This is the external message indicator, which, if supplied,
-        must be a valid 6 letter string of key wheel settings. If not supplied,
-        one will be generated randomly.
-    sys_ind - This is the system indicator, which must be a string of length
-        1 in the range 'A'-'Z', inclusive. If None, one is chosen at random.
-
-    It is an error to supply both key_list and key_list_ind options. Supply one
-    or the other. A ProcedureError will be raised if both are supplied or if
-    neither are supplied.
+    The procedure can be configured with an optional M-209, and optional key
+    list to be used for the day. If the M-209 is not supplied, one will be
+    created internally. Before an encrypt() operation can be performed, a key
+    list must be supplied. This can be done at construction time or via the
+    set_key_list() method.
 
     """
-    # Ensure we have a key list indicator and there is no ambiguity:
+    def __init__(self, m_209=None, key_list=None):
+        self.m_209 = m_209 if m_209 else M209()
 
-    if (key_list and key_list_ind) or (not key_list and not key_list_ind):
-        raise ProcedureError("encrypt requires either key_list or key_list_ind but not both")
+        if key_list:
+            self.set_key_list(key_list)
 
-    # Setup M209 machine if necessary
+    def set_key_list(self, key_list):
+        """Use the supplied key list for all future encrypt operations.
 
-    if key_list:
-        m_209.set_drum_lugs(key_list.lugs)
-        m_209.set_all_pins(key_list.pin_list)
-        key_list_ind = key_list.indicator
+        Configure the M209 with the key list parameters.
 
-    if len(key_list_ind) != 2:
-        raise ProcedureError("key list indicator must be two letters")
+        """
+        if len(key_list.indicator) != 2:
+            raise ProcedureError("invalid key list indicator")
 
-    m_209.letter_counter = 0
+        self.key_list = key_list
+        self.m_209.set_drum_lugs(key_list.lugs)
+        self.m_209.set_all_pins(key_list.pin_list)
 
-    # Set key wheels to external message indicator
-    if ext_msg_ind:
-        try:
-            m_209.set_key_wheels(ext_msg_ind)
-        except M209Error as ex:
-            raise M209Error("invalid external message indicator {} - {}".format(
-                ext_msg_ind, ex))
-    else:
-        ext_msg_ind = m_209.set_random_key_wheels()
+    def encrypt(self, plaintext, group=True, spaces=True, ext_msg_ind=None, sys_ind=None):
+        """Encrypts a plaintext message using standard procedure. The encrypted text
+        with the required message indicators are returned as one string.
 
-    # Ensure we have a valid system indicator
-    if sys_ind:
-        if sys_ind not in M209_ALPHABET_SET:
-            raise ProcedureError("invalid system indicator {}".format(sys_ind))
-    else:
-        sys_ind = random.choice(M209_ALPHABET_LIST)
+        The encrypt function accepts these parameters:
 
-    # Generate internal message indicator
+        plaintext - Input string of text to be encrypted
+        group - If True, the resulting encrypted text will be grouped into 5-letter
+            groups with a space between each group. If False, no spaces will be
+            present in the output.
+        spaces - If True, space characters in the input plaintext will automatically
+            be replaced with 'Z' characters before encrypting.
+        ext_msg_ind - This is the external message indicator, which, if supplied,
+            must be a valid 6 letter string of key wheel settings. If not supplied,
+            one will be generated randomly.
+        sys_ind - This is the system indicator, which must be a string of length
+            1 in the range 'A'-'Z', inclusive. If None, one is chosen at random.
 
-    int_msg_ind = m_209.encrypt(sys_ind * 12, group=False, spaces=False)
+        A ProcedureError will be raised if the procedure does not have a key
+        list to work with.
 
-    # Set wheels to internal message indicator from left to right. We must skip
-    # letters that aren't valid for a given key wheel.
-    it = iter(int_msg_ind)
-    n = 0
-    while n != 6:
-        try:
-            m_209.set_key_wheel(n, next(it))
-        except KeyWheelError:
-            pass
-        except StopIteration:
-            assert False, "Ran out of letters building internal message indicator"
+        """
+        # Ensure we have a key list indicator and there is no ambiguity:
+
+        if not self.key_list:
+            raise ProcedureError("encrypt requires a key list")
+
+        self.m_209.letter_counter = 0
+
+        # Set key wheels to external message indicator
+        if ext_msg_ind:
+            try:
+                self.m_209.set_key_wheels(ext_msg_ind)
+            except M209Error as ex:
+                raise M209Error("invalid external message indicator {} - {}".format(
+                    ext_msg_ind, ex))
         else:
-            n += 1
+            ext_msg_ind = self.m_209.set_random_key_wheels()
 
-    # Now encipher the message on the M209
-    ciphertext = m_209.encrypt(plaintext, group=group, spaces=spaces)
+        # Ensure we have a valid system indicator
+        if sys_ind:
+            if sys_ind not in M209_ALPHABET_SET:
+                raise ProcedureError("invalid system indicator {}".format(sys_ind))
+        else:
+            sys_ind = random.choice(M209_ALPHABET_LIST)
 
-    # If we are grouping, and the final group in the ciphertext has less than
-    # 5 letters, pad with X's to make a complete group:
-    if group:
-        total_len = len(ciphertext)
-        num_groups = total_len // 5
-        num_spaces = num_groups - 1 if num_groups >= 2 else 0
-        x_count = 5 - (total_len - num_spaces) % 5
-        if 0 < x_count < 5:
-            ciphertext = ciphertext + 'X' * x_count
+        # Generate internal message indicator
 
-    # Add the message indicators to pad each end of the message
+        int_msg_ind = self.m_209.encrypt(sys_ind * 12, group=False)
 
-    pad1 = sys_ind * 2 + ext_msg_ind[:3]
-    pad2 = ext_msg_ind[3:] + key_list_ind
+        # Set wheels to internal message indicator from left to right. We must skip
+        # letters that aren't valid for a given key wheel.
+        it = iter(int_msg_ind)
+        n = 0
+        while n != 6:
+            try:
+                self.m_209.set_key_wheel(n, next(it))
+            except KeyWheelError:
+                pass
+            except StopIteration:
+                assert False, "Ran out of letters building internal message indicator"
+            else:
+                n += 1
 
-    msg_parts = [pad1, pad2, ciphertext, pad1, pad2]
+        # Now encipher the message on the M209
+        ciphertext = self.m_209.encrypt(plaintext, group=group, spaces=spaces)
 
-    # Assemble the final message; group if requested
-    sep = ' ' if group else ''
-    return sep.join(msg_parts)
+        # If we are grouping, and the final group in the ciphertext has less than
+        # 5 letters, pad with X's to make a complete group:
+        if group:
+            total_len = len(ciphertext)
+            num_groups = total_len // 5
+            num_spaces = num_groups - 1 if num_groups >= 2 else 0
+            x_count = 5 - (total_len - num_spaces) % 5
+            if 0 < x_count < 5:
+                ciphertext = ciphertext + 'X' * x_count
+
+        # Add the message indicators to pad each end of the message
+
+        pad1 = sys_ind * 2 + ext_msg_ind[:3]
+        pad2 = ext_msg_ind[3:] + self.key_list.indicator
+
+        msg_parts = [pad1, pad2, ciphertext, pad1, pad2]
+
+        # Assemble the final message; group if requested
+        sep = ' ' if group else ''
+        return sep.join(msg_parts)
--- a/m209/tests/test_procedure.py	Fri Jun 07 22:30:33 2013 -0500
+++ b/m209/tests/test_procedure.py	Sat Jun 08 12:55:07 2013 -0500
@@ -7,8 +7,7 @@
 import unittest
 
 from ..keylist import KeyList
-from ..procedure import encrypt
-from ..converter import M209
+from ..procedure import StdEncryptProcedure
 
 
 PLAINTEXT = 'ATTACK AT DAWN'
@@ -29,11 +28,8 @@
                         'ACDFHIMN'
                     ],
                     letter_check = 'TNMYS CRMKK UHLKW LDQHM RQOLW R')
-        self.m_209 = M209()
+        self.proc = StdEncryptProcedure(key_list=self.fm)
 
     def test_encrypt(self):
-
-        result = encrypt(self.m_209, PLAINTEXT, key_list=self.fm,
-                ext_msg_ind='ABCDEF', sys_ind='G')
-
+        result = self.proc.encrypt(PLAINTEXT, ext_msg_ind='ABCDEF', sys_ind='G')
         self.assertEqual(result, CIPHERTEXT)