Enigma

class Rotor:
    def __init__(self, wiring):
        self.wiring = wiring
        self.forward = dict()
        self.backward = dict()
        for i in range(0, 26):
            self.forward[chr(ord("a") + i)] = list(wiring)[
                i
            ]
            self.backward[list(wiring)[i]] = chr(
                ord("a") + i
            )

    def deepcopy(self):
        return Rotor(self.wiring)

    def map(self):
        return lambda ch: self.forward[ch]

    def invmap(self):
        return lambda ch: self.backward[ch]

    def reset(self):
        self.__init__(self.wiring)

    def rotate(self, i=1):
        self.newWiring = self.wiring[i:] + self.wiring[:i]
        self.__init__(self.newWiring)


class Reflector:
    def __init__(self, cfg):
        self.perm = dict()
        self.wiring = cfg
        for i in range(0, 13):
            self.perm[list(cfg)[2 * i]] = list(cfg)[
                2 * i + 1
            ]
            self.perm[list(cfg)[2 * i + 1]] = list(cfg)[
                2 * i
            ]

    def map(self):
        return lambda ch: self.perm[ch]


class PlugBoard:
    def __init__(self, cfg=""):
        self.perm = dict()
        self.wiring = cfg
        for i in range(0, 26):
            self.perm[chr(ord("a") + i)] = chr(ord("a") + i)
        for i in range(0, len(cfg) // 2):
            self.perm[list(cfg)[2 * i]] = list(cfg)[
                2 * i + 1
            ]
            self.perm[list(cfg)[2 * i + 1]] = list(cfg)[
                2 * i
            ]

    def map(self):
        return lambda ch: self.perm[ch]


class Enigma:
    def __init__(self, RotList, Refl, PBoard=PlugBoard()):
        self.rotors = [x.deepcopy() for x in RotList]
        self.reflector = Refl
        self.PlugBoard = PBoard
        self.revo = [0 for x in RotList]

    def setKey(self, PosList):
        self.reset()
        self.revo = [ord(x) - ord("a") for x in PosList]
        for x in zip(self.rotors, self.revo):
            x[0].rotate(x[1])

    def operate(self):
        for x in self.rotors:
            print(x.wiring)
        print("\n")
        self.rotors[0].rotate()
        self.revo[0] += 1
        self.revo[0] %= 26
        for x in range(1, len(self.rotors) - 1):
            if self.revo[x - 1] == 0:
                self.rotors[x].rotate()
                self.revo[x] += 1
                self.revo[x] %= 26

    def reset(self):
        for x in zip(self.rotors, self.revo):
            x[0].rotate(26 - x[1])
        self.revo = [0 for x in self.rotors]

    def encryptChar(self):
        def encr(msg):
            if not (msg.isalpha()):
                return msg
            x = msg.lower()
            x = self.PlugBoard.map()(x)
            for i in self.rotors:
                x = i.map()(x)
            x = self.reflector.map()(x)
            for i in self.rotors[::-1]:
                x = i.invmap()(x)
            x = self.PlugBoard.map()(x)
            self.operate()
            return x

        return encr

    def encode(self):
        return lambda x: "".join(
            self.encryptChar()(c) for c in list(x)
        )


ROT0 = Rotor("zaqrfvbgtyhnmjuiklopcdewsx")
ROT1 = Rotor("qwertyuiopasdfghjklzxcvbnm")
ROT2 = Rotor("qazxswedcvfrtgbnhyujmkiolp")
REF = Reflector("acsxzdqewrfvbgtyhnmjuiklop")
PBOARD = PlugBoard("acdefxoplmnuyt")


ENCMACHINE = Enigma([ROT0, ROT1, ROT2], REF, PBOARD)
DECMACHINE = Enigma([ROT0, ROT1, ROT2], REF, PBOARD)
ENCMACHINE.setKey(["c", "i", "l"])
DECMACHINE.setKey(["c", "i", "l"])


test0 = ENCMACHINE.encode()("a" * 3000)
test1 = DECMACHINE.encode()(test0)