# HG changeset patch
# User Brian Neal <bgneal@gmail.com>
# Date 1370205071 18000
# Node ID c1db39916ec9a5f86a5fab204ca15869aba610ef
# Parent  328790e5bc5d90fe215f07cc291153b736b52d68
Created a decrypt function.

diff -r 328790e5bc5d -r c1db39916ec9 m209/converter.py
--- a/m209/converter.py	Sun Jun 02 13:57:45 2013 -0500
+++ b/m209/converter.py	Sun Jun 02 15:31:11 2013 -0500
@@ -13,8 +13,8 @@
 from .key_wheel import KeyWheel
 from .drum import Drum
 
-ALLOWED_PLAINTEXT = set(string.ascii_uppercase)
-CIPHER = list(reversed(string.ascii_uppercase))
+M209_ALPHABET = set(string.ascii_uppercase)
+CIPHER_TABLE = list(reversed(string.ascii_uppercase))
 
 
 class M209:
@@ -90,31 +90,82 @@
             drum = Drum(lug_list)
         self.drum = drum
 
-    def encrypt(self, plaintext, group=True):
+    def encrypt(self, plaintext, group=True, spaces=True):
+        """Performs an encrypt operation on the given plaintext and returns
+        the ciphertext as a string.
 
+        If group is True, the resulting string will be grouped into 5-letter
+        groups, separated by spaces.
+
+        If spaces is True, space characters in the input plaintext will
+        automatically be treated as 'Z' characters. Otherwise spaces in the
+        plaintext will raise an M209Error exception.
+
+        """
         ciphertext = []
         for p in plaintext:
-            if p not in ALLOWED_PLAINTEXT:
-                raise M209Error("Illegal input to encrypt(): {}".format(p))
-
-            pins = [kw.is_effective() for kw in self.key_wheels]
-            count = self.drum.rotate(pins)
-            c = CIPHER[(ord(p) - ord('A') - count) % 26]
-            ciphertext.append(c)
-
-            for kw in self.key_wheels:
-                kw.rotate()
-
-            self.letter_counter += 1
+            if p == ' ' and spaces:
+                p = 'Z'
+            ciphertext.append(self._cipher(p))
 
         if group:
             s = ' '.join(''.join(ciphertext[i:i+5]) for i in range(0,
                 len(ciphertext), 5))
         else:
             s = ''.join(ciphertext)
+        return s
 
+    def decrypt(self, ciphertext, spaces=True, z_sub=True):
+        """Performs a decrypt operation on the given ciphertext and returns the
+        plaintext as a string.
+
+        If spaces is True, spaces will be allowed in the input ciphertext and
+        ignored. Otherwise space characters will raise an M209Error exception.
+        This is useful if the input ciphertext is in 5-letter groups, separated
+        by spaces, for example.
+
+        If z_sub is True, 'Z' characters in the output plaintext will be
+        replaced by space characters, just like an actual M-209. If z_sub is
+        False, no such substitution will occur.
+
+        """
+        plaintext = []
+        for c in ciphertext:
+            if c == ' ' and spaces:
+                continue
+            plaintext.append(self._cipher(c))
+
+        if z_sub:
+            s = ''.join(s if s != 'Z' else ' ' for s in plaintext)
+        else:
+            s = ''.join(plaintext)
         return s
 
+    def _cipher(self, c):
+        """Simulate a cipher operation on the device:
+        The input letter is read and checked for validity.
+        The guide arm positions are calculated from the current effective pins
+        on the key wheels.
+        The drum is rotated against the guide arms to produce a drum count.
+        The key wheels are rotated one step.
+        The letter counter is incremented.
+        The output letter is computed from the input letter, the drum count, and
+        the internal substitution table.
+
+        """
+        if c not in M209_ALPHABET:
+            raise M209Error("Illegal char: {}".format(c))
+
+        pins = [kw.is_effective() for kw in self.key_wheels]
+        count = self.drum.rotate(pins)
+
+        for kw in self.key_wheels:
+            kw.rotate()
+
+        self.letter_counter += 1
+
+        return CIPHER_TABLE[(ord(c) - ord('A') - count) % 26]
+
 
 if __name__ == '__main__':
 
@@ -129,5 +180,14 @@
     pt = 'A' * 26
     print(m209.encrypt(pt))
 
+    ct = 'OZGPK AFVAJ JYRZW LRJEG MOVLU M'
+    m209 = M209()
+    m209.set_drum_lugs('1-0 2-0*4 0-3 0-4*3 0-5*3 0-6*11 2-5 2-6 3-4 4-5')
+    m209.set_pins(0, 'BFJKLOSTUWXZ')
+    m209.set_pins(1, 'ABDJKLMORTUV')
+    m209.set_pins(2, 'EHJKNPQRSX')
+    m209.set_pins(3, 'ABCHIJLMPQR')
+    m209.set_pins(4, 'BCDGJLNOPQS')
+    m209.set_pins(5, 'AEFHIJP')
+    print(m209.decrypt(ct))
 
-