Updated after tests with real Enigma machine

[cipher-training.git] / enigma.py

diff --git a/enigma.py b/enigma.py
index 8c72f22bf558887842061541b64bfbb7918139f2..a0056e4aae61d34725d63ee82846d29f455859f2 100644 (file)
--- a/enigma.py
+++ b/enigma.py
@@ -182,7 +182,11 @@ class SimpleWheel(LetterTransformer):
             return object.__getattribute__(self, name)
     
     def set_position(self, position):
-        self.position = ord(position) - ord('a')
+        if isinstance(position, str):
+            # self.position = ord(position) - ord('a')
+            self.position = pos(position)
+        else:
+            self.position = position
     
     def forward(self, letter):
         if letter in string.ascii_lowercase:
@@ -234,8 +238,12 @@ class Wheel(SimpleWheel):
             return object.__getattribute__(self, name)
 
     def set_position(self, position):
-        self.position = (pos(position) - self.ring_setting + 1) % 26
-        self.peg_positions = [(pos(p) - pos(position)) % 26  for p in self.ring_peg_letters]
+        if isinstance(position, str):
+            self.position = (pos(position) - self.ring_setting + 1) % 26
+        else:
+            self.position = (position - self.ring_setting) % 26
+        # self.peg_positions = [(pos(p) - pos(position)) % 26  for p in self.ring_peg_letters]
+        self.peg_positions = [(pos(p) - (self.position + self.ring_setting - 1)) % 26  for p in self.ring_peg_letters]
         
     def advance(self):
         super(Wheel, self).advance()