X-Git-Url: https://git.njae.me.uk/?a=blobdiff_plain;f=enigma.py;fp=enigma.py;h=a0056e4aae61d34725d63ee82846d29f455859f2;hb=cc215dec21447bca7c2fb8016df255ab1aac512a;hp=8c72f22bf558887842061541b64bfbb7918139f2;hpb=a718663ac5643f2ee292a6b5784aa4d4ea95ee17;p=cipher-tools.git

diff --git a/enigma.py b/enigma.py
index 8c72f22..a0056e4 100644
--- 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()