# -*- coding: cp1252 -*-
from random import shuffle,randint,choice
from copy import copy
alphabet=range(0,26)

def shift(l, n): # Method to rotate arrays/cogs
    return l[n:] + l[:n]

class cog: # Simple substitution cipher for each cog
    def create(self):
        self.transformation=copy(alphabet)
        shuffle(self.transformation)
        return
    def passthrough(self,i):
        return self.transformation[i]
    def passthroughrev(self,i):
        return self.transformation.index(i)
    def rotate(self):
        self.transformation=shift(self.transformation, 1)
    def setcog(self,a):
        self.transformation=a

class enigma: # Enigma class
    def __init__(self, nocogs,printspecialchars):
        self.printspecialchars=printspecialchars
        self.nocogs=nocogs
        self.cogs=[]
        self.oCogs=[] # Create backup of original cog positions for reset

        for i in range(0,self.nocogs): # Create cogs
            self.cogs.append(cog())
            self.cogs[i].create()
            self.oCogs.append(self.cogs[i].transformation)

        # Create reflector
        refabet=copy(alphabet)
        self.reflector=copy(alphabet)
        while len(refabet)>0:
            a=choice(refabet)
            refabet.remove(a)
            b=choice(refabet)
            refabet.remove(b)
            self.reflector[a]=b
            self.reflector[b]=a

    def print_setup(self): # To print the enigma setup for debugging/replication
        print "Enigma Setup:\nCogs: ",self.nocogs,"\nCog arrangement:"
        for i in range(0,self.nocogs):
            print self.cogs[i].transformation
        print "Reflector arrangement:\n",self.reflector,"\n"

    def reset(self):
        for i in range(0,self.nocogs):
            self.cogs[i].setcog(self.oCogs[i])

    def encode(self,text):
        ln=0
        ciphertext=""
        for l in text.lower():
            num=ord(l)%97
            if (num>25 or num<0):
                if (self.printspecialchars): # readability
                    ciphertext+=l
                else:
                    pass # security
            else:
                ln+=1
                for i in range(0,self.nocogs): # Move thru cogs forward...
                    num=self.cogs[i].passthrough(num)

                num=self.reflector[num] # Pass thru reflector

                for i in range(0,self.nocogs): # Move back thru cogs...
                    num=self.cogs[self.nocogs-i-1].passthroughrev(num)
                ciphertext+=""+chr(97+num) # add encrypted letter to ciphertext

                for i in range(0,self.nocogs): # Rotate cogs...
                    if ( ln % ((i*6)+1) == 0 ): # in a ticker clock style
                        self.cogs[i].rotate()
        return ciphertext

plaintext=raw_input("What do you wanna encrypt ? ")

x=enigma(4,True)
#x.print_setup()

print "Plaintext:\n"+plaintext+"\n"
ciphertext=x.encode(plaintext)
print "Ciphertext:\n"+ciphertext+"\n"

# To proove that encoding and decoding are symmetrical
# we reset the enigma to starting conditions and enter
# the ciphertext, and get out the plaintext
x.reset()
ciphertext2 = raw_input("What do you wanna decode ? ")
decodeplaintext=x.encode(ciphertext2)
print "Plaintext:\n"+decodeplaintext+"\n"