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)