""" enigma.py 2006-09-14 CJ Odenthal """ import operator from string import ascii_uppercase as UPPER from copy import deepcopy from permutations import Permutation,cycleForm,permForm identity = Permutation(UPPER) def char2int(key): try: c = key.upper() return UPPER.index(c) except: return key class Rotor(Permutation): def __init__(self,perm,ringstellung='A',key='A',turnover=''): Permutation.__init__(self,perm) self.turnover = turnover self.setRing(ringstellung) self.setKey(key) def setRing(self,A): r = char2int(A) cycles = [[(x+r) % 26 for x in cycle] for cycle in self.C] perm = permForm(cycles) self.P = perm self.C = cycleForm(perm) def setKey(self,A): k = char2int(A) cycles = [[(x-k) % 26 for x in cycle] for cycle in self.C] perm = permForm(cycles) self.P = perm self.C = cycleForm(perm) self.key = UPPER[k] def step(self): cycles = [[(x-1) % 26 for x in cycle] for cycle in self.C] perm = permForm(cycles) self.P = perm self.C = cycleForm(perm) k = char2int(self.key) k = (k+1) % 26 self.key = UPPER[k] def trigger(self): return self.key in self.turnover class Enigma(object): def __init__(self, config, wheels = None, # 3 regular wheels thin = None, # 0 or 1 thin wheel UKW = None, # Reflector ringstellung = None, # Ring settings grundstellung = None, # Ground setting stecker = None, # Stecker board ): """ config is a dictionary with the following keys: 'wheels' (aka rotors) A dictionary 'thin' (thin rotors) A dictionary (for 4 wheel enigmas) 'UKWs' (aka reflectors) A dictionary 'ETW' (aka QWERTZU) A 'Permutation' 'stecker' (True or False) Is there a stecker? 'dblstep' (True or False) Double stepping of middle rotors? 'UKWring' (True of False) Ring setting for UKW? 'UKWstep' (True of False) Does UKW step? The wheel list should be given from slowest to fastest: W_L, W_M, W_R. The ringstellung and grundstellung are given as strings, e.g. 'ACX'. Any valid options not given during initialization are prompted for. """ # Set the entry wheel. self.ETW = config['ETW'] # Get the wheels. if not wheels: wheels = raw_input("Choose 3 wheels (slow middle fast) from %s: " % config['wheels'].keys()) else: wheels = str(wheels).replace(' ','') wheels = map(int,list(wheels)) self.wheelnumbers = wheels stdwheels = [deepcopy(config['wheels'][i]) for i in wheels] # Get the thin wheel if needed. if config['thin']: if not thin: thin = raw_input("Choose a thin wheel from %s: " % config['thin'].keys()) self.thin = [deepcopy(config['thin'][thin.upper()])] else: self.thin = [] # Get the reflector. if len(config['UKWs']) > 1: if not UKW: UKW = raw_input("Choose a reflector from %s: " % config['UKWs'].keys()) self.UKW = [deepcopy(config['UKWs'][UKW.upper()])] else: self.UKW = [deepcopy(config['UKWs']['A'])] # Assemble all the wheels. self.basewheels = self.UKW + self.thin + stdwheels # Get the stecker if needed. if config['stecker']: if not stecker: stecker = raw_input("Choose an even number of letters to swap in pairs for your stecker board: ") else: stecker = identity self.setStecker(stecker) # Get the number of rings. self.nR = len(self.basewheels) if not config['UKWring']: self.nR -= 1 # Get the ringstellung. if not ringstellung: ringstellung = raw_input("Choose a %s character 'ringstellung': " % (self.nR,)) ringstellung = ringstellung.replace(' ','').upper() self.__setRings(ringstellung) # Get the grundstellung. if not grundstellung: grundstellung = raw_input("Choose a %s character 'grundstellung': " % (self.nR,)) grundstellung = grundstellung.replace(' ','').upper() self.setKeys(grundstellung) # Set up the correct stepping function if config['dblstep']: self.step = self.__step2 else: # Do some tweaking, then set the step function. self.turnSets = [set(W.turnover) for W in self.wheels[1-self.nR:]] self.holdSets = [set() for k in range(self.nR-1)] self.step = self.__step1 def __repr__(self): return repr(self.permutation) def __str__(self): return str(self.permutation) def __call__(self,text): return "".join(map(self.encipher,text)) def setStecker(self,stecker): try: steck = stecker.replace(' ','').upper() steck = map(UPPER.index,steck) N = len(steck) stecker = [(steck[2*k],steck[2*k+1]) for k in range(N/2)] except AttributeError: pass try: self.stecker = Permutation(stecker) except IndexError: self.stecker = identity except TypeError: self.stecker = stecker self.ES = self.ETW * self.stecker def __setRings(self,ringstellung): # Convert ring setting letters to numbers. ringstellung = map(char2int,ringstellung) # All ring settings default to 'A'. nS = len(ringstellung) ringstellung = (0,)*(self.nR - nS) + tuple(ringstellung) for k,R in enumerate(self.basewheels[-self.nR:]): R.setRing(ringstellung[k]) def setKeys(self,grundstellung): # Get base set of wheels in 'AAA...' configuration. self.wheels = deepcopy(self.basewheels) # Convert key setting letters to numbers. grundstellung = map(char2int,grundstellung) # All key settings default to 'A'. nK = len(grundstellung) grundstellung = (0,)*(self.nR - nK) + tuple(grundstellung) for k,R in enumerate(self.wheels[-self.nR:]): R.setKey(grundstellung[k]) # A new key setting requires recalculation of the permutation. self.setPermutation() def setPermutation(self): self.permutation = (reduce(operator.lshift,self.wheels) << self.ES) def keys(self): keyList = [R.key for R in self.wheels[-self.nR:]] return "".join(keyList) def turnovers(self): for R in self.wheels[-self.nR:]: print R.turnover, def __step1(self): # This implements the 'single stepping' used in the Abwehr enigma. # Something mysterious is going on here with the 'holdSets'. # I'm blindly hoping it isn't screwing things up. nR = self.nR for i in range(1-nR, 0): k = self.wheels[i].key if k in self.turnSets[i]: if i<-1: self.turnSets[i].discard(k) self.holdSets[i].add(k) self.wheels[i-1].step() if i > 1-nR: self.turnSets[i-1].update(self.holdSets[i-1]) self.holdSets[i-1] = set() self.wheels[-1].step() self.setPermutation() def __step2(self): # This implements the famous 'double stepping' nR = self.nR F = nR - 1 motion = [0]*nR for i in range(-F,0): if self.wheels[i].trigger(): motion[i-1] = 1 motion[i] = 1 break motion[-1] = 1 for i in range(nR): if motion[i]: self.wheels[i-nR].step() self.setPermutation() def encipher(self,c): if c.upper() in UPPER: self.step() return self.permutation.encipher(c) else: return c