Bits of tinkering

[cipher-training.git] / cipher.py

diff --git a/cipher.py b/cipher.py
index 0fb6dc255d4ce42656783b1e28f7c2caa379a0b0..b53ae24184d09ba2b6bcf304c0de8a2f0077f318 100644 (file)
--- a/cipher.py
+++ b/cipher.py
@@ -212,7 +212,7 @@ def affine_decipher_letter(letter, multiplier=1, adder=0, one_based=True):
         if one_based: cipher_number += 1
         plaintext_number = ( 
             modular_division_table[multiplier]
-                                  [(cipher_number - adder) % 26] )
+                                  [(cipher_number - adder) % 26])
         if one_based: plaintext_number -= 1
         return chr(plaintext_number % 26 + alphabet_start) 
     else:
@@ -540,7 +540,10 @@ class PocketEnigma(object):
         else:
             self.validate_wheel_spec(wheel)
             self.make_wheel_map(wheel)
-        self.position = ord(position) - ord('a')
+        if position in string.ascii_lowercase:
+            self.position = ord(position) - ord('a')
+        else:
+            self.position = position
 
     def make_wheel_map(self, wheel_spec):
         """Expands a wheel specification from a list of letter-letter pairs
@@ -563,7 +566,7 @@ class PocketEnigma(object):
         >>> pe.validate_wheel_spec([])
         Traceback (most recent call last):
             ...
-        ValueError: Wheel specification has 0 pairs, require 13
+        ValueError: Wheel specification has 0 pairs, requires 13
         >>> pe.validate_wheel_spec([('a', 'b', 'c')]*13)
         Traceback (most recent call last):
             ...
@@ -574,7 +577,7 @@ class PocketEnigma(object):
         ValueError: Wheel specification does not contain 26 letters
         """
         if len(wheel_spec) != 13:
-            raise ValueError("Wheel specification has {} pairs, require 13".
+            raise ValueError("Wheel specification has {} pairs, requires 13".
                 format(len(wheel_spec)))
         for p in wheel_spec:
             if len(p) != 2:
@@ -604,8 +607,16 @@ class PocketEnigma(object):
         5
         >>> ''.join([pe.lookup(l) for l in string.ascii_lowercase])
         'udhbfejcpgmokrliwntsayqzvx'
+        >>> pe.lookup('A')
+        ''
         """
-        return chr((self.wheel_map[(ord(letter) - ord('a') - self.position) % 26] + self.position) % 26 + ord('a'))
+        if letter in string.ascii_lowercase:
+            return chr(
+                (self.wheel_map[(ord(letter) - ord('a') - self.position) % 26] + 
+                    self.position) % 26 + 
+                ord('a'))
+        else:
+            return ''
 
     def advance(self):
         """Advances the wheel one position.
@@ -618,7 +629,7 @@ class PocketEnigma(object):
         self.position = (self.position + 1) % 26
         return self.position
 
-    def encipher(self, message):
+    def encipher(self, message, starting_position=None):
         """Enciphers a whole message.
 
         >>> pe.set_position('f')
@@ -629,7 +640,11 @@ class PocketEnigma(object):
         5
         >>> pe.encipher('kjsglcjoqc')
         'helloworld'
+        >>> pe.encipher('helloworld', starting_position = 'x')
+        'egrekthnnf'
         """
+        if starting_position:
+            self.set_position(starting_position)
         transformed = ''
         for l in message:
             transformed += self.encipher_letter(l)