Python Forum

Full Version: Enigma Machine Working - Optimise
You're currently viewing a stripped down version of our content. View the full version with proper formatting.
Hi Everyone,
This is my first forum post so apologies for any mistakes.
I recently wrote this code for the enigma machine, it is not well designed or laid out as I am a beginner.
Any feedback on the code or the layout would be much appreciated, or even just the concept.
Thanks Michael

alphabet = ['A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z']
rot = [
	[[15, 4, 25, 20, 14, 7, 23, 18, 2, 21, 5, 12, 19, 1, 6, 11, 17, 8, 13, 16, 9, 22, 0, 24, 3, 10],0],
	[[25, 14, 20, 4, 18, 24, 3, 10, 5, 22, 15, 2, 8, 16, 23, 7, 12, 21, 1, 11, 6, 13, 9, 17, 0, 19],0],
	[[4, 7, 17, 21, 23, 6, 0, 14, 1, 16, 20, 18, 8, 12, 25, 5, 11, 24, 13, 22, 10, 19, 15, 3, 9, 2],0],
	[[8, 12, 4, 19, 2, 6, 5, 17, 0, 24, 18, 16, 1, 25, 23, 22, 11, 7, 10, 3, 21, 20, 15, 14, 9, 13],0],
	[[16, 22, 4, 17, 19, 25, 20, 8, 14, 0, 18, 3, 5, 6, 7, 9, 10, 15, 24, 23, 2, 21, 1, 13, 12, 11],0]
		 ]
rotors = [
	[[15, 4, 25, 20, 14, 7, 23, 18, 2, 21, 5, 12, 19, 1, 6, 11, 17, 8, 13, 16, 9, 22, 0, 24, 3, 10],0],
	[[25, 14, 20, 4, 18, 24, 3, 10, 5, 22, 15, 2, 8, 16, 23, 7, 12, 21, 1, 11, 6, 13, 9, 17, 0, 19],0],
	[[4, 7, 17, 21, 23, 6, 0, 14, 1, 16, 20, 18, 8, 12, 25, 5, 11, 24, 13, 22, 10, 19, 15, 3, 9, 2],0],
	[[8, 12, 4, 19, 2, 6, 5, 17, 0, 24, 18, 16, 1, 25, 23, 22, 11, 7, 10, 3, 21, 20, 15, 14, 9, 13],0],
	[[16, 22, 4, 17, 19, 25, 20, 8, 14, 0, 18, 3, 5, 6, 7, 9, 10, 15, 24, 23, 2, 21, 1, 13, 12, 11],0]
		 ]
reflector = [
	('A', 'E'), ('B', 'J'), ('C', 'M'), ('D', 'Z'), ('H', 'X'), ('I', 'V'),
	('L', 'F'), ('Q', 'O'), ('R', 'N'), ('T', 'S'), ('U', 'P'),
	('W', 'K'), ('Y', 'G'),
			]
#Creates the rotors for the enigma machine

def createPlugboard():
	global Plugboard
	Plugboard = []
	print('You are now requested to enter 10 pairs of letters,\neach pair must consist of unique letters and no letter can appear twice!')
	while len(Plugboard) != 10:
		temp = set(alphabet)
		for i in Plugboard:
			i = set(i)
			temp = temp - i
		print('The unchoosen letters are:',' '.join(temp))
		flag = True
		print('Pair -',len(Plugboard)+1)
		a , b = (input('Letter 1- ').upper() , input('Letter 2- ').upper())
		for i in Plugboard:
			if a in i or b in i:
				flag = False
				print('Letter allready choosen')
				#checks to see if letter allready in plugboard
		if a not in alphabet or b not in alphabet:
			flag = False
			print('Please enter a single letter')
			#Checks to see if letter is an individual charector and an actual letter
		elif a == b:
			flag = False
			print('Please enter different letters')
		if flag:
			Plugboard.append((a,b))
			#Finally adds the pair to the plugboard
		if Plugboard:
			print('Current plugboard is:')
			for i in Plugboard: print('	',i[0],'and',i[1])
		else:
			print('Current plugboard is:\nEmpty')

def selectRotors():
	global RotorInUse
	RotorInUse = []
	print('You will now select the order you would like the rotors in,\nthere are five rotors (1,2,3,4,5).')
	availableRotor = ['1','2','3','4','5']
	while len(RotorInUse) < 3:
		choice = input()
		#user chooses the rotors
		if choice not in availableRotor:
			print('Please choose an available rotor!')
			#checks to see if rotor in use
		else:
			RotorInUse.append(int(choice))
			#if not in use adds selected rotor
		print('Your Current Rotor order is:',RotorInUse)
	for i in range(len(RotorInUse)):
		RotorInUse[i] -=1

def RotorSettings():
	global rotorPositions , RotorInUse , rotors , rot
	print(rot)
	print('You will now have to set the rotor positions')
	print(rotors)
	rotors = rot
	print(rotors)
	#resets the rotors to default position
	rotorPositions = []
	desiredSetting = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', '14', '15', '16', '17', '18', '19', '20', '21', '22', '23', '24', '25']
	while len(rotorPositions) <3:
		flag = True
		setting = input('Please enter the desired position 0-25 for rotor',RotorInUse[len(rotorPositions)])
		if not(setting not in desiredSetting or setting in rotorPositions):
			rotorPositions.append(int(setting))
	for i in range(len(RotorInUse)):
		for j in range(rotorPositions[i]):
			rotors[RotorInUse[i]][0].append(rotors[RotorInUse[i]][0].pop(0))
			rotors[RotorInUse[i]][1] += 1

def ThroughRotor(letter):
	global rotors , RotorInUse
	r1 = rotors[RotorInUse[0]][0]
	r2 = rotors[RotorInUse[1]][0]
	r3 = rotors[RotorInUse[2]][0]
	#asigns easy variable names
	letter = alphabet.index(letter)
	letter = r3[r2[r1[letter]]]
	#first time through the rotors
	letter = alphabet[letter]
	#back into letter
	letter = reflect(letter)
	#through the reflector, switching letter
	letter = alphabet.index(letter)
	#back into a number
	letter = r3.index(letter)
	letter = r2.index(letter)
	letter = r1.index(letter)
	#through the rotors one final time
	letter = alphabet[letter]
	rotate()
	#adds one to the rotor possition
	return letter
	#returns the letter

def plugs(letter):
	global Plugboard
	p = Plugboard
	for i in range(len(p)):
		if letter in p[i]:
			letter = p[i][(1-p[i].index(letter))]
	return letter

def reflect(letter):
	global reflector
	r = reflector
	for i in range(len(r)):
		if letter in r[i]:
			letter = r[i][(1-r[i].index(letter))]
	return letter
	
	
def encode():
	string = input('Please enter the string you would like to encode or decode\n- ')
	encoded = ''
	for i in string:
		if i.upper() in alphabet:
			temp = plugs(i.upper())
			temp = ThroughRotor(temp)
			temp = plugs(temp)
			encoded += temp
		else:
			encoded += i
	return encoded
			
		
def rotate():
	global rotors , RotorInUse , rotorPositions
	r1 = RotorInUse[0]
	r2 = RotorInUse[1]
	r3 = RotorInUse[2]
	rotors[r1][1] += 1
	rotors[r1][0].append(rotors[r1][0].pop(0))
	if rotors[r1][1] == 26:
		rotors[r1][1] = 0
		rotors[r2][1] += 1
		rotors[r2][0].append(rotors[r2][0].pop(0))
		if rotors[r2][1] == 26:
			rotors[r2][1] = 0
			rotors[r3][1] += 1
			rotors[r3][0].append(rotors[r3][0].pop(0))
			if rotors[r3][1] == 26:
				rotors[r3][1] = 0
	rotorPositions[0] = rotors[r1][1]
	rotorPositions[1] = rotors[r2][1]
	rotorPositions[2] = rotors[r3][1]
def RotorMenu():
	global RotorInUse
	while True:
		choice = input('Please choose one of the following options;\n	1. New Rotors in machine\n	2. New Rotor positions\n	3. Exit\n- ')
		if choice =='2':
			if RotorInUse:	
				RotorSettings()			
			else:
				print('Please select which rotors you want first')
		elif choice == '1':
			selectRotors()
		elif choice == '3':
			break

def PlugboardMenu():			
	global Plugboard
	while True:
		choice = input('Please choose one of the following options;\n	1. Create completely new Plugboard pairs\n	2. Change a plugboard Pair\n	3. Exit\n- ')
		if choice =='1':
			createPlugboard()
		elif choice == '2':
			pass
			'''if Plugboard:
				ChangePair()
			else:
				print('Please Complete the plugboard before you attempt to change it!')'''
		elif choice == '3':
			break
	
def ChangePair():
	global Plugboard
	flag = True
	input('These are the pairs that you currently have set up!\n Please press enter to continue - ')
	for i in Plugboard:
		print(i[0],',',i[1])
	while flag:
		choice = input('Please enter a letter in the pair you wish to change,\n and then confirm that the selection is correct.\n If you wish to exit, simply type "exit"\n -').upper()
		if choice == 'exit':
			return
		for i in Plugboard:
			if choice in i:
				print(i)
				cont = input('Please confirm that the above pair is the one you wish to change,\nIf it is not then enter "n"\nIf it is correct enter "y"\n -').lower()
				if cont == 'y':
					flag = False
					choice = Plugboard.index(i)
	flag = True
	while flag:
		Plugboard.pop(choice)
		print('Please enter your new pair;')
		a , b = input('Letter 1: ').upper() , input('Letter 2: ').upper()
		if a not in alphabet or b not in alphabet:
			print('Please enter a single letter')
		else:
			for i in Plugboard:
				if a in i or b in i:
					print('Please enter letters that are not in use')
				else:
					Plugboard.append((a,b))
	
	
	
RotorInUse = [0,1,2]
Plugboard = [('Q', 'W'), ('E', 'R'), ('T', 'Y'), ('U', 'I'), ('O', 'P'), ('A', 'S'), ('D', 'F'), ('G', 'H'), ('J', 'K'), ('L', 'Z')]
rotorPositions = [0,0,0]
				
while True:
	print('\nPlease choose one of the following options;\n	1. Encode / Decode Message\n	2. Change the rotor settings\n	3. Change the plugboard\n	4. Complete Fresh Setup')
	print('Here is a list of the itmes you currently have set up and their settings:\nRotor Positions:',rotorPositions,'\nRotors in machine:',[i+1 for i in RotorInUse],'\nPlugboard:',Plugboard)
	choice = input('- ')
	
	if choice == '1':
		if RotorInUse and rotorPositions and Plugboard:	
			print(encode())
			input('Press enter to continue!')
		else:
			print('Please set up machine!')
	elif choice == '2':
		RotorMenu()
	elif choice == '3':
		PlugboardMenu()
	elif choice == '4':
		createPlugboard()
		selectRotors()
		RotorSettings()
	
	
The first and most obvious thing is to stop using global. You want to be clear about what is going on in your functions. You do that with parameters and return values. See the function tutorial on how to do that.

The second thing is that this application really begs for classes. Take selectRotors. You have availableRotors and RotorsInUse, but neither of those contain rotors. They contain numbers pointing to rotors. So later on rather than just using the rotors, you have to use those numbers to find the rotors. With classes, availableRotors could just be a list of Rotor objects. When they select one, you could just remove it from availableRotors and put it in RotorsInUse, which would then be full of rotors. Check out the class basics tutorial for information on how to do that.
Mike, I think your code and the concept is cool. But how do I decode a previously coded phrase? I copied and pasted the coded phrase under option 1 (Encode / Decode) and it Encodes again.....

Jeff
I worked on a version of enigma, never finished, but I did create GUI for the patchboard here: https://python-forum.io/Thread-Enigma-Ma...9#pid31469 which you might be interested in.

And also this site: http://users.telenet.be/d.rijmenants/en/enigmasim.htm
I did a suite of classic cryptography once. I think it went up to the Enigma machine. I'm not sure where it is any more. It might be on the old site.