Mercurial > public > enigma
changeset 14:0fe1c4a11bad
Created a factory function to make it easier to create Enigmas.
author | Brian Neal <bgneal@gmail.com> |
---|---|
date | Sun, 27 May 2012 15:05:00 -0500 |
parents | 3fbdc7005075 |
children | 7ea7da689dbd |
files | enigma/machine.py enigma/main.py enigma/tests/test_enigma.py |
diffstat | 3 files changed, 44 insertions(+), 28 deletions(-) [+] |
line wrap: on
line diff
--- a/enigma/machine.py Sun May 27 13:57:13 2012 -0500 +++ b/enigma/machine.py Sun May 27 15:05:00 2012 -0500 @@ -8,6 +8,10 @@ """ import string +from .rotors.factory import create_rotor, create_reflector +from .plugboard import Plugboard + + class EnigmaError(Exception): pass @@ -46,6 +50,40 @@ self.reflector = reflector self.plugboard = plugboard + @classmethod + def from_key_sheet(cls, rotors, ring_settings=(0, 0, 0), reflector='B', + plugboard_settings=''): + """Convenience function to build an EnigmaMachine from the data as you + might find it on a key sheet: + + rotors: a list of strings naming the rotors from left to right; + e.g. ["I", "III", "IV"] + + ring_settings: an iterable of integers representing the ring settings to + be applied to the rotors in the rotors list + + reflector: a string that names the reflector to use + + plugboard: a string of plugboard settings as you might find on a key + sheet; e.g. 'PO ML IU KJ NH YT GB VF RE DC' + + """ + # validate inputs + num_rotors = len(rotors) + if num_rotors not in (3, 4): + raise EnigmaError("invalid rotors list size") + + if num_rotors != len(ring_settings): + raise EnigmaError("please provide %d ring settings" % num_rotors) + + # assemble the machine + rotor_list = [create_rotor(r[0], r[1]) for r in zip(rotors, ring_settings)] + + return cls(rotor_list, + create_reflector(reflector), + Plugboard(plugboard_settings)) + + def set_display(self, val): """Sets the rotor operator windows to 'val'.
--- a/enigma/main.py Sun May 27 13:57:13 2012 -0500 +++ b/enigma/main.py Sun May 27 15:05:00 2012 -0500 @@ -2,21 +2,14 @@ # This file is part of Py-Enigma, the Enigma Machine simulation. # Py-Enigma is released under the MIT License (see License.txt). +from .machine import EnigmaMachine -from rotors.factory import create_rotor -from rotors.factory import create_reflector -from machine import EnigmaMachine def main(): - rotors = [] - rotors.append(create_rotor('I')) - rotors.append(create_rotor('II')) - rotors.append(create_rotor('III')) - - reflector = create_reflector('B') - - machine = EnigmaMachine(rotors=rotors, reflector=reflector) + machine = EnigmaMachine.from_key_sheet( + rotors=['I', 'II', 'III'], + reflector='B') machine.set_display('AAA') cipher_text = machine.process_text('AAAAA')
--- a/enigma/tests/test_enigma.py Sun May 27 13:57:13 2012 -0500 +++ b/enigma/tests/test_enigma.py Sun May 27 15:05:00 2012 -0500 @@ -6,9 +6,7 @@ import unittest -from ..rotors.factory import create_rotor, create_reflector from ..machine import EnigmaMachine -from ..plugboard import Plugboard class SteppingTestCase(unittest.TestCase): @@ -21,12 +19,7 @@ # This example taken from # http://users.telenet.be/d.rijmenants/en/enigmatech.htm # in the section on "The Stepping Mechanism." - rotors = [] - rotors.append(create_rotor("III")) - rotors.append(create_rotor("II")) - rotors.append(create_rotor("I")) - - m = EnigmaMachine(rotors, create_reflector('B'), Plugboard()) + m = EnigmaMachine.from_key_sheet(rotors=['III', 'II', 'I']) m.set_display('KDO') @@ -43,15 +36,7 @@ CIPHER_TEXT = 'BDZGO' def setUp(self): - rotors = [] - rotors.append(create_rotor('I')) - rotors.append(create_rotor('II')) - rotors.append(create_rotor('III')) - - reflector = create_reflector('B') - - self.machine = EnigmaMachine(rotors=rotors, reflector=reflector, - plugboard=Plugboard()) + self.machine = EnigmaMachine.from_key_sheet(rotors=['I', 'II', 'III']) self.machine.set_display('AAA') def test_simple_encrypt(self):