From 9a530c46fc9b380f4195a99cd460d67092408259 Mon Sep 17 00:00:00 2001 From: Neil Smith Date: Thu, 10 Jul 2014 22:24:53 +0100 Subject: [PATCH] Tweaks to pocket enigma --- cipher.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/cipher.py b/cipher.py index 0fb6dc2..713db29 100644 --- a/cipher.py +++ b/cipher.py @@ -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 @@ -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) -- 2.34.1