Enigma Machine Simulation in Python

Question

Background and Example

This code simulates the Enigma machine, minus the plugboard. Here's some test code that illustrates how the machine's construction and use:

>>> r1 = Rotor("VEADTQRWUFZNLHYPXOGKJIMCSB", 1)
>>> r2 = Rotor("WNYPVJXTOAMQIZKSRFUHGCEDBL", 2)
>>> r3 = Rotor("DJYPKQNOZLMGIHFETRVCBXSWAU", 3)
>>> reflector = Reflector("EJMZALYXVBWFCRQUONTSPIKHGD")
>>> machine = Machine([r1, r2, r3], reflector)
>>> x = machine.encipher("ATTACK AT DAWN")
>>> machine.decipher(x)
'ATTACK AT DAWN'

Rotors

The Rotor class is pretty simple. A Rotor knows how to rotate itself, and provides methods for navigating connections with the adjacent circuits through the encipher and decipher methods.

class Rotor:
 """
 Models a 'rotor' in an Enigma machine
 Rotor("BCDA", 1) means that A->B, B->C, C->D, D->A and the rotor has been
 rotated once from ABCD (the clear text character 'B' is facing the user)
 Args:
 mappings (string) encipherings for the machine's alphabet.
 offset (int) the starting position of the rotor
 """
 def __init__(self, mappings, offset=0):
 self.initial_offset = offset
 self.reset()
 self.forward_mappings = dict(zip(self.alphabet, mappings))
 self.reverse_mappings = dict(zip(mappings, self.alphabet))
 def reset(self):
 """
 Helper to re-initialize the rotor to its initial configuration
 Returns: void
 """
 self.alphabet = Machine.ALPHABET
 self.rotate(self.initial_offset)
 self.rotations = 1
 def rotate(self, offset=1):
 """
 Rotates the rotor the given number of characters
 Args: offset (int) how many turns to make
 Returns: void
 """
 for _ in range(offset):
 self.alphabet = self.alphabet[1:] + self.alphabet[0]
 self.rotations = offset
 def encipher(self, character):
 """
 Gets the cipher text mapping of a plain text character
 Args: character (char)
 Returns: char
 """
 return self.forward_mappings[character]
 def decipher(self, character):
 """
 Gets the plain text mapping of a cipher text character
 Args: character (char)
 Returns: char
 """
 return self.reverse_mappings[character]

Reflector

Pretty straightforward. A Reflector can reflect a character and is used to put the input back through machine's rotors.

class Reflector:
 """
 Models a 'reflector' in the Enigma machine. Reflector("CDAB")
 means that A->C, C->A, D->B, B->D
 Args: mappings (string) bijective map representing the reflection
 of a character
 """
 def __init__(self, mappings):
 self.mappings = dict(zip(Machine.ALPHABET, mappings))
 for x in self.mappings:
 y = self.mappings[x]
 if x != self.mappings[y]:
 raise ValueError("Mapping for {0} and {1} is invalid".format(x, y))
 def reflect(self, character):
 """
 Returns the reflection of the input character
 Args: character (char)
 Returns: char
 """
 return self.mappings[character]

Machine

This class exposes the encipher and decipher methods. Most of the enciphering is done through the helper function encipher_character.

class Machine:
 """
 Models an Enigma machine (https://en.wikipedia.org/wiki/Enigma_machine)
 Args:
 rotors (list[Rotor]) the configured rotors
 reflector (Reflector) to use
 """
 ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
 def __init__(self, rotors, reflector):
 self.rotors = rotors
 self.reflector = reflector
 def encipher(self, text):
 """
 Encipher the given input
 Args: text (string) plain text to encode
 Returns: string
 """
 return "".join((self.encipher_character(x) for x in text.upper()))
 def decipher(self, text):
 """
 Deccipher the given input
 Args: text (string) cipher text to decode
 Returns: string
 """
 for rotor in self.rotors:
 rotor.reset()
 return self.encipher(text)
 def encipher_character(self, x):
 """
 Runs a character through the machine's cipher algorithm
 1. If x is not in the known character set, don't encipher it
 2. For each of the rotors, determine the character in contact with x.
 Determine the enciphering for that character, and use it as the next
 letter to pass through to the next rotor in the machine's sequence
 3. Once we get to the reflector, get the reflection and repeat the above
 in reverse
 4. Rotate the first rotor, and check if any other rotor should be rotated
 5. Return the character at the terminating contact position as the input
 character's enciphering
 Args: x (char) the character to encode
 Returns: char
 """
 if x not in Machine.ALPHABET:
 return x
 # compute the contact position of the first rotor and machine's input
 contact_index = Machine.ALPHABET.index(x)
 # propagate contact right
 for rotor in self.rotors:
 contact_letter = rotor.alphabet[contact_index]
 x = rotor.encipher(contact_letter)
 contact_index = rotor.alphabet.index(x)
 # reflect and compute the starting contact position with the right rotor
 contact_letter = Machine.ALPHABET[contact_index]
 x = self.reflector.reflect(contact_letter)
 contact_index = Machine.ALPHABET.index(x)
 # propagate contact left
 for rotor in reversed(self.rotors):
 contact_letter = rotor.alphabet[contact_index]
 x = rotor.decipher(contact_letter)
 contact_index = rotor.alphabet.index(x)
 # rotate the first rotor and anything else that needs it
 self.rotors[0].rotate()
 for index in range(1, len(self.rotors)):
 rotor = self.rotors[index]
 turn_frequency = len(Machine.ALPHABET)*index
 if self.rotors[index-1].rotations % turn_frequency == 0:
 rotor.rotate()
 # finally 'light' the output bulb
 return Machine.ALPHABET[contact_index]

Improvements

I'm wondering how the encipher algorithm might be implemented more cleanly (in particular, how the code might be better distributed across the Rotor and Reflector classes). Overall comments on the style and documentation would be much appreciated, too.

FYI: development for this project has been moved to: https://github.com/gjdanis/enigma

Answer 1

Your Rotor.rotate can be simplified to

def rotate(self, offset=1):
 self.rotations = offset
 self.alphabet = self.alphabet[offset:] + self.alphabet[:offset]

This saves having to do a costly list addition offset times.

---

Different commonly used ASCII character classes are included in the string module. You can use string.ascii_uppercase for the uppercase alphabet.

---

join can take a generator expression directly, so you can get rid of one set of parenthesis in Machine.encipher:

return "".join(self.encipher_character(x) for x in text.upper())

---

It is better to ask forgiveness than permission. One place where you can use this is the check whether the character to encode is in the character set. Just use try..except in Machine.encipher_character:

 # compute the contact position of the first rotor and machine's input
 try:
 contact_index = Machine.ALPHABET.index(x)
 except ValueError:
 return x

This way you avoid having to go through the list more often than necessary (the edge case Z will iterate through the alphabet twice, once to see if it is there and once to actually get the index).

---

In the same function you have the two blocks # propagate contact right and # propagate contact left. The comments already suggest that this would be a perfect place to make them a function. They also only differ by whether or not the rotors are traversed in reverse or not and whether to use rotor.encipher or rotor.decipher. Make a method Machine.rotate_rotors:

def rotate_rotors(self, left=False):
 """propagate contact right or left"""
 iter_direction = reversed if left else iter
 for rotor in iter_direction(self.rotors):
 contact_letter = rotor.alphabet[self.contact_index]
 func = rotor.decipher if left else rotor.encipher
 self.contact_index = rotor.alphabet.index(func(contact_letter))

I would also add the reflector rotating into a method:

def rotate_reflector(self):
 """reflect and compute the starting contact position with the right rotor"""
 contact_letter = Machine.ALPHABET[self.contact_index]
 x = self.reflector.reflect(contact_letter)
 self.contact_index = Machine.ALPHABET.index(x)

You can then use these like this:

self.contact_index = Machine.ALPHABET.index(x)
self.rotate_rotors()
self.rotate_reflector()
self.rotate_rotors(left=True)

I made contact_index a property of the class now, this way we don't have to pass in the contact_index every time and return it. I also made your comments into docstrings.

• rookie: Thanks! Can you say more about how you mean to put the code for the reflector in rotate_rotors? I was thinking about making this function an inner function to encipher_character. What do you think about that?
• Graipher: @rookie Regarding the further simplification, I'll think about it some more. But I would let the function be a method of Machine, this way if you ever want to rotate all rotors(/reflectors) there is a method for that that is not hidden within another method.
• Graipher: @rookie I can't think of a way that is not too complex. Just put it into its own method, Machine.rotate_reflector(self, contact_index).

Answer 2

You say:

This class exposes the encipher and decipher methods. Most of the enciphering is done through the helper function encipher_character.

Therefore, I would assume that encipher_character should be private since it is not exposed and a helper function (which are often private). You can change this to either __encipher_character or _encipher_character to indicate this. You can read more about private methods here (it explains also why there are two options).

• rookie: Yep -- I could do that, or I could move the upper() call inside encipher_character and then leave the method public. Then, encipher and encipher_character would have the same behavior for a single character.