from Enigma.Rotor import Rotor
from Enigma.Reflector import Reflector
from Enigma.Plugboard import Plugboard

class Enigma:
    
    def __init__(self , rotors = [ Rotor(0,"IC") , Rotor(0,"IIC") , Rotor(0,"IIIC") ] , plugboard = Plugboard() , reflector = Reflector("A")): 
        self.rotors = rotors
        for i in range(len(rotors)):
            if i + 1 < len(rotors):
                rotors[i].on("Sidereal", lambda *args: rotors[i+1].step())
        self.Plugboard = plugboard;
        self.Reflector = reflector;

    def encrypt(self,data):
        data = data.upper().replace(" ","");
        string = "";
        for char in data:
            string += self.each(char,True);
        return string;

    def decrypt(self,data):
        data = data.upper();
        string = "";
        for char in data:
            string += self.each(char,False);
        return string;

    def each(self,char,flag):
        self.rotors[0].step()
        output = self.Plugboard.get(char)
        for rotor in self.rotors:
            if flag:
                output = rotor.scramble(output)
            else:
                output = rotor.unscramble(output)
        output = self.Reflector.get(output)
        for rotor in self.rotors[::-1]:
            if flag:
                output = rotor.scramble(output)    
            else:
                output = rotor.unscramble(output)  
        return self.Plugboard.get(output);

[((7, 100, 7, 111), 'Enigma.Plugboard.Plugboard', 'Plugboard', ({}, {}), '()', False, 'from Enigma.Plugboard import Plugboard\n'), ((7, 126, 7, 140), 'Enigma.Reflector.Reflector', 'Reflector', ({(7, 136, 7, 139): '"""A"""'}, {}), "('A')", False, 'from Enigma.Reflector import Reflector\n'), ((7, 35, 7, 48), 'Enigma.Rotor.Rotor', 'Rotor', ({(7, 41, 7, 42): '(0)', (7, 43, 7, 47): '"""IC"""'}, {}), "(0, 'IC')", False, 'from Enigma.Rotor import Rotor\n'), ((7, 51, 7, 65), 'Enigma.Rotor.Rotor', 'Rotor', ({(7, 57, 7, 58): '(0)', (7, 59, 7, 64): '"""IIC"""'}, {}), "(0, 'IIC')", False, 'from Enigma.Rotor import Rotor\n'), ((7, 68, 7, 83), 'Enigma.Rotor.Rotor', 'Rotor', ({(7, 74, 7, 75): '(0)', (7, 76, 7, 82): '"""IIIC"""'}, {}), "(0, 'IIIC')", False, 'from Enigma.Rotor import Rotor\n')]