'''
Polyalphabetic encryption, decryption
Create N alphabets, then for character at position i, use
alphabet k for substitution, where k = i mod N.
For simplicity, we will not encrypt characters not in alphabet

Author: Tim Pierson, Dartmouth CS55, Winter 2021
				Based Wenliang Du https://github.com/kevin-w-du/BookCode/blob/master/Encryption/poly_sub_cipher.py
Usage: python3 polyalphabet.py <plaintext filename>
'''

import sys, random, string
N = 1000 #number of ciphers to create
alphabet = "abcdefghijklmnopqrstuvwxyz"

'''
Create N mappings of alphabet characters.  Each character in the alphabet is 
represented by another character, like in the monoalphabetic solution, but now there
are N such mappings.  Each mapping is a key.  The position of a character in the key
corresponds to the alphabet character (e.g. the first character in the key is 'a', the second is 'b', ...)

Input: 
	N -- number of keys to make
Output: 
	List of N strings of scrambled alphabet characters as encryption keys
'''
def make_keys(N):
	keys = []
	for i in range(N):
		#scramble order of letters in alphabet (e.g. a key)
		mapping = random.sample(alphabet,len(alphabet))
		#add to keys list
		keys.append(''.join(mapping))
	return keys

'''
Encrypt and text using the N encryption keys created by the 
make_key routine.

Input: 
	keys -- list of N strings to used to encrypt plaintext
	plaintext -- text to encrypt
Output:
	text encrypted using keys
'''
def encrypt(keys,plaintext):
	#set up blank ciphertext with one character for each character in plaintext
	ciphertext = [None] * len(plaintext)
	#loop over each plaintext character, one at a time
	for i in range(0, len(plaintext)):
		#if plaintext character not in alphabet, just encode with that character
		#(e.g., encode a comma as a comma, a period as a period, space as a space, etc)
		plain_character = plaintext[i] #get character
		if plain_character in alphabet:
			#character is in alphabet (e.g, is a-z)
			#get ith key from list of keys (wrap around if i > N)
			key = keys[i%N]
			#determine position of character in alphabet (e.g., a at 0, b at 1, ...)
			position = alphabet.find(plain_character) 
			#pick the character at position from the key
			ciphertext[i] = key[position]
		else:
			#character is not in alphabet, just encode as is
			ciphertext[i] = plain_character
	#return ciphertext as string
	return ''.join(ciphertext)


'''
Decrypt ciphertext using the N encryption keys created by the 
make_key routine.

Input: 
	keys -- list of N strings to used to decrypt ciphertext
	ciphertext -- text to decrypt
Output:
	decrypted ciphertext now as plaintext
'''
def decrypt(keys,ciphertext):
	#set up blank plaintext with one character for each character in ciphertext
	plaintext = [None] * len(ciphertext)
	#loop over each ciphertext character
	for i in range(0, len(ciphertext)):
		#if ciphertext character not in alphabet, just decode with that character
		#(e.g., decode a comma as a comma, a period as a period, space as a space, etc)
		cipher_character = ciphertext[i]
		if cipher_character in alphabet:
			#character is in alphabet (e.g, is a-z)
			#get ith key from list of keys (wrap around if i > N)
			key = keys[i%N]
			#determine position of character in key 
			position = key.find(cipher_character)
			#pick the character at position from the alphabet
			plaintext[i] = alphabet[position]
		else:
			#character is not in alphabet, just decode as is
			plaintext[i] = cipher_character
	#return plaintext as string
	return ''.join(plaintext)

if __name__ == '__main__':
	#usage check: make sure we got a filename to read
	if len(sys.argv) != 2:
		print("Usage: python3 polyalphabet.py <plaintext filename>")
		print("Did you provide a filename?")
		exit()

	#read plaintext file 
	f = open(sys.argv[1], "r")
	plaintext = f.read()
	f.close()
	print("Plain text:")
	print(plaintext)

	#create N encryption keys
	keys = make_keys(N)
	
	#encrypt plaintext using N keys
	print("\nCiphertext:")
	ciphertext = encrypt(keys,plaintext)
	print(ciphertext)

	#decrypt the ciphertext using same N keys
	print("\nDecrypted ciphertext")
	plain = decrypt(keys,ciphertext)
	print(plain)