7
\$\begingroup\$

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

\$\endgroup\$
4
\$\begingroup\$

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.

\$\endgroup\$
  • \$\begingroup\$ 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? \$\endgroup\$ – rookie Aug 24 '16 at 14:16
  • \$\begingroup\$ @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. \$\endgroup\$ – Graipher Aug 25 '16 at 6:21
  • \$\begingroup\$ @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). \$\endgroup\$ – Graipher Aug 25 '16 at 16:12
1
\$\begingroup\$

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).

\$\endgroup\$
  • \$\begingroup\$ 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. \$\endgroup\$ – rookie Aug 24 '16 at 3:11

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.