Enigma

1.py

from enigma import *

machine = Enigma()
machine.setRotorPositions('A','A','A','A')
machine.setPlugboard({12:3,2:21,13:16,6:14,5:17,18:25,20:8,24:19,11:10,9:0,15:7,22:1,4:23})
result = machine.decrypt('gnzbhnlxjikohyymbfhmitoicdossllwinhybzr',True)

print(result)

2.2.py

from enigma import *

message = 'attackatdawn'
answer = 'tgrmaosivkmz'
plugboard = {i:i for i in range(26)}

machine = Enigma()
machine.setRotorPositions('W','Q','V','D')

messageNum = machine.toNumbers(message)
answerNum = machine.toNumbers(answer)

def increment(n):
	return (n + 1) % 26

def testInference(plaintextMessage,encodedMessage,startingIndex,or0,or1,or2,or3,originalPlugboard,emptyPlugs):
	plugboard = originalPlugboard.copy()
	r0,r1,r2,r3 = or0,or1,or2,or3
	
	machine.setPlugboard(plugboard);
	
	for i in range(startingIndex,len(plaintextMessage)):
		inputLetter = plaintextMessage[i];

		#	if we have fully tested the inference and no failures have occurred yet, return true
		if inputLetter == originalPlugboard[inputLetter] and inputLetter in emptyPlugs:
			break

		#	test one more letter
		expectedLetter = machine.toNumber(encodedMessage[i])
		outputLetter = machine.rotorizeLetter(machine.toNumber(inputLetter),r0,r1,r2,r3)

		#	if there is a discrepancy between the rotor output and the expected letter in the crib, try to add a new jumper
		if outputLetter != expectedLetter:
			#	if we must add a plug to a hole that is already filled or must remain empty (based on previous assumptions), return false
			if (outputLetter != originalPlugboard[outputLetter]) or (expectedLetter != originalPlugboard[expectedLetter]) or (outputLetter not in emptyPlugs) or (expectedLetter not in emptyPlugs):
				return -1

			#	adding the required jumper does not cause a contradiction, so continue
			plugboard = machine.jump((outputLetter,expectedLetter),plugboard.copy());
			
			machine.setPlugboard(plugboard)
			#	else: via rotors and current jumpers, the letter we got back from the machine matches the expected output... great!
		
		#	save the rotor settings for the next iteration
#		rotorSettings = output.getRotorSettings();

	#	we have successfully tested the inference as far as possible, so return
	originalPlugboard = plugboard
	or0,or1,or2,or3 = r0,r1,r2,r3

	return i

def infer(plaintextMessage,encodedMessage,startingIndex,r0,r1,r2,r3,originalPlugboard,originalEmptyPlugs):
	plugboards = []
	#	if we've infered past the end of the string, terminate with the current plugboard
	if startingIndex >= len(plaintextMessage):
		plugboards.append(originalPlugboard)
		return plugboards
	#	start guessing from the current letter
	guessFrom = plaintextMessage[startingIndex]
	for guessTo in range(26):
		#	if we are trying to guess on a plug that already has a jumper in it, skip
		if guessTo != originalPlugboard[guessTo]:
			continue

		newPlugboard = originalPlugboard.copy()
		newEmptyPlugs = originalEmptyPlugs.copy()
	
		if guessFrom == guessTo:
			#	if we are guessing that a plug has a jumper to itself, mark it as "intentionally left blank" :)
	 	   newEmptyPlugs.add(guessFrom)
		else:
			#	try connecting up our guess
	 	   newPlugboard = machine.jump((guessFrom,guessTo),newPlugboard.copy())

		#	with our new assumption, make as many conclusions as possible and detect contradictions
		endingIndex = testInference(plaintextMessage,encodedMessage,startingIndex,r0,r1,r2,r3,newPlugboard,newEmptyPlugs)

		if endingIndex > -1:
			#	our assumption seems valid so far, so keep making assumptions
			childPlugboards = infer(plaintextMessage,encodedMessage,endingIndex,r0,r1,r2,r3,newPlugboard,newEmptyPlugs)

			#	merge the results from the child calls to the current results
			plugboards.reserve(childPlugboards.size() + plugboards.size());
			for i in range(len(childPlugboards)):
				plugboards.append(childPlugboards[i])

	return plugboards

result = infer(message,answer,0,machine.r0,machine.r1,machine.r2,machine.r3,plugboard,set())
print(result)

2.py

from enigma import *

message = 'attackatdawn'
answer = 'tgrmaosivkmz'
plugboard = {i:i for i in range(26)}
used = set()

machine = Enigma()
machine.setRotorPositions('W','Q','V','D')

messageNum = machine.toNumbers(message)
answerNum = machine.toNumbers(answer)

def increment(n):
	return (n + 1) % 26

#def guess(x,y,i):
#	if x == y:
#		y = increment(x)
#	board[x] = y
#	board[y] = x
#	machine.setPlugboard(board)
#	result = machine.encrypt(message)
#	if result[i] == message[i]:
#		used.add(number)
#		print(result)
#	if (x in used and not board[x] == y) or (y in used and not board[y] == x):
#		guess(x,increment(x))
#	if x == 26:
#		guess(increment(x),0)

def guess(index,plugboard,r0,r1,r2,r3):
	if index >= len(message):
		return plugboard
	board = machine.jump((messageNum[index],messageNum[index] + 1),plugboard.copy())
	machine.setPlugboard(board)
	
	newLetter = machine.toNumber(message[index])
	newLetter = machine.plug(newLetter)
	newLetter = machine.rotorizeLetter(newLetter,r0,r1,r2,r3)
	
	board = machine.jump((newLetter,answerNum[index]),board)
	
	newLetter = machine.plug(newLetter)
	
	r0 = (r0 + 1) % 26
	if r0 == 0:
		r1 = (r1 + 1) % 26
		if r1 == 0:
			r2 = (r2 + 1) % 26
			if r2 == 0:
				r3 = (r3 + 1) % 26
	if message[:index] == answer[:index]:
		return guess(index + 1,board,r0,r1,r2,r3)
	else:
		return plugboard

attempt = guess(messageNum[0],plugboard,machine.r0,machine.r1,machine.r2,machine.r3)

machine.setPlugboard(attempt)
print(machine.encrypt(message))

enigma.py

import math

class Enigma:
	def __init__(self):
		self.log = 1
		
		self.prob = [7.25,1.25,3.5,4.25,12.75,3,2,3.5,7.75,0.25,0.5,3.75,2.75,7.75,7.5,2.0,0.5,8.5,6,9.25,3,1.5,1.5,0.5,2.25,0.25]
#		self.setPlugboard(({12,3},{2,21},{13,16},{6,14},{5,17},{18,25},{20,8},{24,19},{11,10},{9,0},{15,7},{22,1},{4,23}))
#		self.setPlugboard([
#			[12,2,13,6,5,18,20,24,11,9,15,22,4],
#			[3,21,16,14,17,25,8,19,10,0,7,1,23]
#		])
		self.setRotorPositions('A','A','A','A')
		self.setPlugboard({i:i for i in range(26)})
		self.setRotors([
			[22,5,1,19,8,25,13,20,21,2,16,4,9,15,14,10,18,7,23,6,12,0,3,24,17,11],
			[7,12,17,6,13,16,15,10,9,0,2,8,21,22,11,20,4,1,5,25,18,24,23,19,3,14],
			[3,17,20,0,7,11,1,5,25,21,6,12,22,2,10,23,15,8,16,24,18,14,13,19,9,4],
			[14,22,19,21,18,10,24,15,9,8,5,12,11,20,0,7,17,16,4,2,13,3,1,25,6,23]
		])
		
		for p in range(len(self.prob)):
			self.prob[p] = math.log(self.prob[p],2)
	
	def setAnswer(self,answer):
		self.answer = ''.join(self.toLetters(answer))
	
	def setPlugboard(self,plugboard):
		self.plugboard = [
			plugboard,
			{plugboard[k]:k for k in plugboard}
		]
	
	def setRotors(self,rotors):
		self.rotors = rotors
		self.rotorsInverse = rotors[0:-1]
		for r,rotor in enumerate(self.rotorsInverse):
			temp = [None] * 26
			for l,letter in enumerate(rotor):
				temp[letter] = l
			self.rotorsInverse[r] = temp
	
	def setRotorPositions(self,r0,r1,r2,r3):
		self.r0 = self.toNumber(r0)
		self.r1 = self.toNumber(r1)
		self.r2 = self.toNumber(r2)
		self.r3 = self.toNumber(r3)
	
	def toNumber(self,letter):
		return ord(letter.lower()) - 97

	def toNumbers(self,letters):
		result = []
		for letter in letters:
			result.append(ord(letter.lower()) - 97)
		return result
	
	def toLetter(self,number):
		return chr(number + 97)
	
	def toLetters(self,numbers):
		result = []
		for number in numbers:
			result.append(chr(number + 97))
		return result
	
	def jump(self,jumper,board):
		board[jumper[0]] = jumper[1]
		board[jumper[1]] = jumper[0]
		return board

	def plug(self,message):
		def swap(letter):
			if letter in self.plugboard[0]:
				return self.plugboard[0][letter]
			else:
				return self.plugboard[1][letter]
#			try:
#				return self.plugboard[0][self.plugboard[1].index(letter)]
#			except:
#				return self.plugboard[1][self.plugboard[0].index(letter)]
		
		result = []
		if type(message) == list:
			for letter in range(len(message)):
				result.append(swap(message[letter]))
		elif type(message) == int:
			result = swap(message)
		return result
	
	def rotorizeLetter(self,letter,r0,r1,r2,r3):
		letter = (self.rotors[0][(letter + r0) % 26] - r0) % 26
		letter = (self.rotors[1][(letter + r1) % 26] - r1) % 26
		letter = (self.rotors[2][(letter + r2) % 26] - r2) % 26

		letter = (self.rotors[3][(letter + r3) % 26] - r3) % 26

		letter = (self.rotorsInverse[2][(letter + r2) % 26] - r2) % 26
		letter = (self.rotorsInverse[1][(letter + r1) % 26] - r1) % 26
		letter = (self.rotorsInverse[0][(letter + r0) % 26] - r0) % 26
		
		return letter

	def rotorize(self,message,r0,r1,r2,r3):
		result = []
		for letter in range(len(message)):
			crypt = self.rotorizeLetter(message[letter],r0,r1,r2,r3)	
			result.append(crypt)
		
			r0 = (r0 + 1) % 26
			if r0 == 0:
				r1 = (r1 + 1) % 26
				if r1 == 0:
					r2 = (r2 + 1) % 26
					if r2 == 0:
						r3 = (r3 + 1) % 26
		return result
	
	def encrypt(self,message,bypass = 'none'):
		answer = self.toNumbers(message)
		if bypass is not 'plugboard':
			answer = self.plug(answer)
		if bypass is not 'rotors':
			answer = self.rotorize(answer,self.r0,self.r1,self.r2,self.r3)
		if bypass is not 'plugboard':
			answer = self.plug(answer)
		self.setAnswer(answer)
		return self.answer
	
	def decrypt(self,message,echo = False):
		plugged = self.plug(self.toNumbers(message))

		for x in range(26 ** 4):
			frequency = 0
	
			attempt = self.rotorize(plugged,self.r0,self.r1,self.r2,self.r3)
			attempt = self.plug(attempt)
	
			for letter in attempt:
				frequency += self.prob[letter]
	
			if frequency > self.log:
				self.log = frequency
				self.setAnswer(attempt)
				if echo:
					print(self.toLetter(self.r0),self.toLetter(self.r1),self.toLetter(self.r2),self.toLetter(self.r3),'|',self.answer,'|',self.log)

			self.r3 = (self.r3 + 1) % 26
			if self.r3 == 0:
				self.r2 = (self.r2 + 1) % 26
				if self.r2 == 0:
					self.r1 = (self.r1 + 1) % 26
					if self.r1 == 0:
						self.r0 = (self.r0 + 1) % 26
		return self.answer