
#Use the included script run.sh to execute gpg_package_handler - execute run using ./run.sh
#the options available to run.sh are shown below
#  debugFlag    - If value is 1 print error/descriptive messages.
#  pathToScript - The path leading to the folder containing this script.
#				  (Used if run.sh is launched outside of folder containing gpg_package_handler.py)

#Things to do
#TODO ensure data is not cached in insecure locations. done in run.sh by making sure gpg has uid of root so sensitive data is handled in RAM and not page files.
#TODO find better way to detect if error occurs while sending key to keyserver.
#TODO complete documentation

#Testing to do
#TODO confirm consistent strategy for error handling
#TODO confirm consistent use of new line within print statements - after message only.
#TODO test with various file types - tested with .png, .docx  and works as expected
#TODO test use with subkeys

#Future / Experimental work
#TODO implement argument passing from command line to allow this utility to be used within scripts
#TODO offline mode - dont query keyserver when decrypting message - allows for key forgery though!

import configparser
import fileinput
import gnupg
import os
import re #regular expression
import subprocess
import sys
import time
from datetime import date
from datetime import datetime
from optparse import OptionParser

#Classes

class KeyArgs:
	def __init__(self, name, email, expiry_date, custom_date=False):
		self.name = name
		self.email = email
		self.expiry_date = expiry_date
		self.custom_date = custom_date

class KeyInfo:
	def __init__(self, fingerprint, key_id=None, uid=None):
		self.fingerprint = fingerprint
		self.key_id = key_id
		self.uid = uid

	def print_info(self):
		print('User ID: %s' % self.uid)
		print('Fingerprint: %s' % self.fingerprint)
		print('Key ID: %s' % self.key_id)
		print()


class SenderArguments:
	# userid (name,email), recipients (recip1, recip2), keytype, keylength, subkeytype, subkeylength, expiry date,
	def __init__(user_id=None, recipient_list=None, key_type=None, key_length=None, subkey_type=None, subkey_length=None, expiry_date=None):
		self.user_id = user_id
		self.recipient_list = recipient_list
		self.key_type = key_type
		self.key_length = key_length
		self.subkey_type = subkey_type
		self.subkey_length = subkey_length
		self.expiry_date = expiry_date

# Global Variables

sender_path = os.getcwd()+ '/data/sender/'
receiver_path = os.getcwd() + '/data/receiver/'
sender_pubkey_path = receiver_path + 'sender_public_keys/'
encrypt_path =  sender_path + 'to_encrypt/'
encrypted_path = receiver_path +  'to_decrypt/'
decrypted_path = receiver_path + 'decrypted/'
default_key_server = 'hkp://pool.sks-keyservers.net'

# Functions

def print_exit(messages=""):
	if(messages != ""):
	    print(messages)
	print("- Exiting Program.")
	print()
	raise SystemExit

def print_keyring_info(secret=False):

	KeyInfos = []
	i = 0

	#If secret is True, only the keypairs that have a private key associated with them are retrieved.
	import_result = gpg.list_keys(secret)
	
	print(" -- Available Keys -- ")
	for fingerprint in import_result.key_map:

		# Do not show expired or revoked keys
		key_info = retrieve_valid_key(fingerprint)
		if(key_info == None):
			continue;
	
		i += 1
		if(i%10 == 0):
			input('Press enter to show next 10 available keys.\n')

		print('%s. %s\n   Fingerprint: %s' % (str(i), key_info.uid, key_info.fingerprint))
		KeyInfos.append(key_info)
		
	print()
	return KeyInfos

def collect_recipient_info():	
	name = input('What is the name of the recipient? : ')
	email = input('What is the email of the recipient? : ')

	return format_uid(name,email)

def reassign_default_keypair():

	# List all keys without including secret keys and associate a number with each key
	# Select which key to assign to default key using input of associated number.
	complete_keypairs_only = True
	key_info_list = print_keyring_info(complete_keypairs_only)
	list_len = len(key_info_list)
	if(len(key_info_list) == 0):
		print('- Unable to find any keypairs within GPG keyring.')
		return 1

	#loop controlling valid selection of default key
	while(True):
		#loop controlling valid entry of a numeric argument
		while(True):
			reply = None
			reply = collect_input(prompt = 'Enter number associated with desired key.\n - ', pattern='[1-9]+')
			print()
			if(reply != None):
				break

		if(int(reply) > list_len):
			print('- Invalid selection. Key does not exist\n')
			continue
		break

	reply = int(reply) - 1
	selected_key = key_info_list[int(reply)]
	
	try:
		set_default_keypair(selected_key.fingerprint)
	except Exception as e:
		print(e)
		raise Exception('Unable to reassign default keypair.\n')
		
		
	print('Successfully set default key to that of %s.\n' % selected_key.uid)
	return 0

def format_uid(name, email):
	return str(name + ' <' + email + '>')

def collect_input(prompt,pattern):
	reply = None
	reply = input(prompt)

	if(re.fullmatch(pattern,reply,flags=0) == None):
		print("Input is incorrect.\n")
		return None

	return reply

#optional arguments to allow for future use of command line arguments with GPG Package Handler
#def prepare_key_arguments(userId=None, keyType=None, keyLength=None, subkeyType=None, subkeyLength=None, expiryDate=None):
def prepare_key_arguments():
	reply = None
	#name,email,date,custom_date
	key_args = KeyArgs(None,None,None,False)

	print('- Creating new keypair.')
	print('Enter credentials to be associated with this key.\n')
	key_args.name = input('Enter Name: ')
	key_args.email = input('Enter Email Address: ')

	print('- Creating expiry date for new default key.\n')
	print('Please provide your input.')

	while(True):
		reply = None
		reply = collect_input(prompt='1. Custom.\n2. 1 Year.\n3. 1 Week.\n4. 1 Day.\n - ',pattern='[1-4]')
		if(reply != None):
			break
	print()

	if(reply != '1'):
		key_args.custom_date = False

	if(reply == '2'):
		key_args.expiry_date = '1y'

	if(reply == '3'):
		key_args.expiry_date = '1w'

	if(reply == '4'):
		key_args.expiry_date = '1d'

	#if a custom date was provided
	if(reply == '1'):
		key_args.custom_date = True

		print('Please provide a custom date in the shown format yyyy-mm-dd.')
		#print('You may also provide a date in a combination of the following formats.')
		#print('Days (1d), Weeks (1w) or Years (1y).')

		while(True):
			#for when the first date is incorrect
			key_args.expiry_date = None
			key_args.expiry_date = input('Please enter date : ')
			print()

			try:
				key_args.expiry_date = time.strptime(key_args.expiry_date ,'%Y-%m-%d')
			except: #Not using exception message as the information provided is adequate
				print('Please enter date using the format yyyy-mm-dd. Ex 2016-12-31')
				continue

			key_args.expiry_date = datetime.fromtimestamp(time.mktime(key_args.expiry_date)).date()
			if(key_args.expiry_date <= date.today()):
				print('The expiry date must be in the future.')
				continue
			break
			#end while loop

	return key_args

def generate_keypair(name, email, date, keyServer):

	PRIMARY_UID = 0

	print("Generating keypair using GnuPG. Please wait...")

	#input_data = gpg.gen_key_input(expire_date=str(date), key_type="DSA", key_length=4096, name_real=name, name_email=email, subkey_type="ELG-E", subkey_length=3072)
	input_data = gpg.gen_key_input(expire_date=str(date), key_type="RSA", key_length=4096, name_real=name, name_email=email)
	try:
		new_keypair = gpg.gen_key(input_data)
		#print(new_keypair)
	except Exception as e:
		print(e)
		raise Exception(' - Error occured while generating keypair with GnuPG.')

	#public key
	#try:
		#public key converted to ascii for support with web clients
	#	ascii_armored_public_key = gpg.export_keys(new_keypair.fingerprint, armor=True)
	#except Exception as e:
	#	print(e)
	#	raise Exception(' - Error creating ASCII armored public key. Key cannot be provided to recipient(s)')

	send_result=gpg.send_keys(keyServer,new_keypair.fingerprint)
	
	#ensuring the key was sent to the key server.	
	errors = str(send_result.__dict__['stderr'])		
	if(errors.find("keyserver send failed") != -1):	
		#If your public key cannot be sent to the keyserver then the sender cannot retreive your public key to send you messages.					
		gpg.delete_keys(new_keypair.fingerprint,secret=True) #deleting the private key
		gpg.delete_keys(new_keypair.fingerprint)  # deleting the public key
		raise Exception('- Error sending key to keysever.') 

	#Retrieve key_id and uid from result of list_keys
    #Since the keypair has just been generated only a single uid can exist for this keypair. 
	key_dict = gpg.list_keys(keys=new_keypair.fingerprint)[0]

	return KeyInfo(key_id=key_dict['keyid'], fingerprint = new_keypair.fingerprint, uid = key_dict['uids'][PRIMARY_UID])

#optional arguments allowing for future use of command line arguments with GPG Package Handler
#def create_default_keypair(userId=None, keyType=None, keyLength=None, subkeyType=None, subkeyLength=None, expiryDate=None, keyServer=None):
def create_default_keypair(keyServer=None):
	key_args = prepare_key_arguments()

	if(keyServer == None):
		keyServer = default_key_server
	
	if(key_args.custom_date == True):
		try:
			key_args.expiry_date = key_args.expiry_date.strftime('%Y-%m-%d')
		except Exception as e:
			err_output='Error while formatting custom date of %s' % key_args.expiry_date
			raise Exception(err_output)
	try:
		key_info = generate_keypair(key_args.name,key_args.email,key_args.expiry_date, keyServer)
	except Exception as e:
		print(e)
		raise Exception('- Error creating keypair.')

	try:
		set_default_keypair(key_info.fingerprint)
		print('- Created keypair successfully set as default key.')
		print('- The default key can be reset from the main menu.')
	except Exception as e:
		gpg.delete_keys(key_info.fingerprint,secret=True) #deleting the private key
		gpg.delete_keys(key_info.fingerprint); #deleting the public key
		print(e)
		raise Exception('Unable to set created key as default key.')

	return key_info

def set_default_keypair(fingerprint):

	gpg_config = None
	gpg_temp_config = None
	def_key_set=0
	def_key_entry='default-key '+ fingerprint + '\n'

	#open gpg.conf at location specified by gpg_path
	gpg_conf_path = gpg_path + "/gpg.conf"
	gpg_temp_path = gpg_path + "/gpg_temp.conf"

	try:

		#Verify the authenticity of the fingerprint by exporting it out of the keyring.
		#If the export fails, this key is not in the keyring and cannot be set as default key.

		if(len(gpg.export_keys(fingerprint)) == 0):
			raise('Key matching provided fingerprint could not be found in keyring.')

		#check if gpg configuration file exists
		#if not create a copy of this file
		#add default key entry to the file

		if(os.path.isfile(gpg_conf_path) == False):
			gpg_config = open(gpg_conf_path, 'w')
			gpg_config.write(def_key_entry)
			gpg_config.close()
		else:
			#modify gpg configuration file by performing these steps
			#create temp config file from current config file.
			#locate default-key entry in the configuration file
			#replace the keyid value with the supplied fingerprint
			#copy contents of temp config to active config

			gpg_config = open(gpg_conf_path, 'r') #read
			gpg_temp_config = open(gpg_temp_path, 'w') #write

			#locate default-key entry in the configuration file
			for line in gpg_config:
				if(line.find('default-key') != -1):
					match = re.match('^(default-key|#default-key)\s([A-Z0-9]{40}|[A-Z0-9]{8})', line, flags=0)
					if(match != None):
						#replace the keyid value with the supplied fingerprint
						line=str(def_key_entry)
						def_key_set=1
				#create temp file from current file, the entire file must be read to allow a full copy to be made
				gpg_temp_config.write(line)

			if(def_key_set == 0):
				gpg_temp_config.write(def_key_entry)

			#Close config file to allow it to be reopened in write IO mode.
			gpg_config.close()
			#close temp file to allow the buffer to flush its contents to file
			gpg_temp_config.close()

			gpg_config = open(gpg_conf_path,'w')
			gpg_temp_config = open(gpg_temp_path, 'r')
			
			#copy contents of temp config to active config
			gpg_config.write(gpg_temp_config.read())

			gpg_config.close()
			gpg_temp_config.close()
			os.remove(gpg_temp_config.name)

	except Exception as e:
			if(gpg_config != None):
				gpg_config.close()

			if(gpg_temp_config != None):
				gpg_temp_config.close()

			print(e)
			raise Exception('Error occured while assigning key to default key.')

def get_default_keypair():

	gpg_conf_path = gpg_path + "/gpg.conf"
	gpg_config = None
	fingerprint = None
	def_key = None	

	#locate default-key entry in the configuration file
	try:

		gpg_config = open(gpg_conf_path, 'r') #read

		for line in gpg_config:
			if(line.find('default') != -1):
				match = re.match('^default-key\s([A-Z0-9]{40}|[A-Z0-9]{8})', line, flags=0)
				if(match != None):
					fingerprint = match.group(1)

					# retrieving key entered in configuration file from GnuPG's keyring.
					def_key = retrieve_valid_key(fingerprint,verbose=True)
					if(def_key != None):
						gpg_config.close()
						return def_key

		#end of for loop
		gpg_config.close()

	except Exception as e:
		if(gpg_config != None):
			gpg_config.close()
		print(e)
		print_exit('- Error retrieving default public key')

	print("A default keypair has not be assigned.\n")
	return None

def retrieve_valid_key(fingerprint="", verbose = False):

	KEY_INDEX=0
	PRIMARY_UID=0
	key_list=None
	key=None

	#Retreiving keys from GnuPG keyring.
	key_list = gpg.list_keys(keys=fingerprint)
	
	#print(key_list)

	if(len(key_list) != 0):
		key = key_list[KEY_INDEX]

	if (key == None):
		if(verbose == True):
			print('A key with the following fingerprint could not be found within GnuPG\'s keyring:\n%s\n' % fingerprint)
		return None

	#If key is expired.
	if (str(key['trust']) == 'e'):
		if(verbose == True):
			print('%s\'s key is expired.\nKey Fingerprint: %s\n' % (key['uids'][PRIMARY_UID],fingerprint))
		return None

    #If key is revoked.
	if (str(key['trust']) == 'r'):
		if(verbose == True):
			print('%s\'s key has been revoked.\nKey Fingerprint: %s\n' % (key['uids'][PRIMARY_UID],fingerprint))
		return None

	#only the first uid is selected as this is the primary uid of the keypair. 
	return KeyInfo(fingerprint,key['keyid'],key['uids'][PRIMARY_UID])

def search_for_key(recipients, key_server = None):

	key_list = []
	recipient_info = []
	i = 0
	reply = None


	if(key_server == None):
		key_server = default_key_server

	#if only a single recipient is to be searched for
	if(type(recipients) is str):
		s = recipients
		recipients = []
		recipients.append(s)

	for recipient in recipients:
		match = 0
		while(True): #start of loop allowing for reattempt of search for specific recipient
			try:
				print('- Searching for %s within %s.\n' % (recipient,key_server))
				key_list = gpg.search_keys(recipient, key_server)
			except Exception as e:
				print(e)
				print('Error while searching for public key of %s.\n' % recipient)
				break

			if(key_list.__dict__['curkey'] != None):
				#ensuring the provided user ID matches the searched user ID
				#cannot compare by fingerprint at this point as the fingerprint is unknown before the search by User ID.		
				for key in key_list:				
					for uid in key['uids']:
						if(uid == recipient):
							match = 1
							recipient_info.append(KeyInfo(fingerprint = key['keyid'],uid=uid))
							break

					if(match == 1):
						break
				#end of loop through key_list

			if(match != 1):
				print('- Unable to find information on recipient %s from %s.\n' % (recipient, key_server))
				prompt = 'Choose retry option below:\n'
				options = [ '1. Search using new recipient info.\n','2. Retry search for ' + recipient + '.\n','3. Abandon search for ' + recipient +'.\n', '4. Return to main menu.\n' ]
				
				choice = create_menu(prompt, options)
				
				print(choice);
				
				if(choice == 1):
					recipients[i] = collect_recipient_info()
					recipient = recipients[i]
					continue

				if(choice == 2):
					continue

				if(choice == 3):
					break
				
				if(choice == 4):
					recipient_info = [];
					break;

			print('- Recipient %s found at %s.\nFingerprint of public key: %s \n' % (recipient, key_server, recipient_info[i].fingerprint))
			i += 1
			break; #end of while loop allowing reattempt of search for specific recipient
		#end of recipients loop
	return recipient_info

def output_to_file(path, public_key):

	pubkey_file = None

	try:
		#Outputting the public key to file
		pubkey_file = open(path, 'w')
		pubkey_file.write(public_key)
		pubkey_file.close()

	except Exception as e:

		if(pubkey_file != None):
			pubkey_file.close()

		print(e)
		raise Exception("File could not be exported")

def encrypt_file(encrypt_path, encrypted_path, recip_fingerprints, signer_fingerprint):

	file = None

	print('- Beginning file encryption.\n')

	try:
		file = open(encrypt_path, 'rb')
		encrypted_path = encrypted_path + '.gpg';
		result = gpg.encrypt_file(file, recip_fingerprints, sign=signer_fingerprint, output=encrypted_path)
	except Exception as e:
		#print(result.data)
		#print(result.ok)
		#print(result.status)

		if(file != None):
			file.close()

		print(e)
		raise Exception("Unable to encrypt file.")

	file.close()
	print('- Successfully encrypted and signed file.\nEncrypted file has been output to %s\n' % encrypted_path)
	
	
#def sign_file(file_path, signer_fingerprint, userid = None):
#
#	signed_data = None
#	file = None

#	print()
#	print('- Signing file located at:\n' + file_path)
#	print('Signing file with private key of ' + userId)
#	file = open(file_path, 'rb')
#
#	try:
#       try clearsign = False
#		signed_data=gpg.sign_file(file, keyid=signer_fingerprint, clearsign=True, output=file_path)
#	except Exception as e:

#		if(file != None):
#			file.close()

#		reprint(e)
#       raise Exception('Error signing file with the private key of ' + userid)
#		raise Exception('Error signing file with the private key.')

#	file.close()
#	return signed_data


#Checking if a keypair's public key has previously been signed by the sender or receiver. 
def check_for_key_signature(key_to_sign_fingerprint, signature_to_find_keyid):
	
	cur_key_id = None	
	RECEIVER_KEY = 0
	KEY_ID = 0
	
	key_list = gpg.list_keys(keys=key_to_sign_fingerprint,sigs=True)

	if(len(key_list) != 0):		
		for sig_info in key_list[RECEIVER_KEY]['sigs']:					
			cur_key_id = str(sig_info[KEY_ID])
			if(cur_key_id == signature_to_find_keyid):					
				return True
	
	return False

#sign_key calls on BASH script sign_key.sh that calls upon gpg's --sign-key command
#key signing is done using the current default key.
def sign_key(fingerprint):

	resources_path = os.getcwd() + "/res/sign_key.sh"
	reply = None
	key_info = retrieve_valid_key(fingerprint,verbose=True)
	
	if(key_info == None):
		return False

	print( "-- Details of key to be signed --\n")
	key_info.print_info()

	while(True):
		reply = collect_input(prompt="Do you want to sign this key\n1. Yes\n2. No\n - ",pattern="1|2")
		print()
		if(reply != None):
			break
	
	if(reply == '1'):
		try:
			retcode = subprocess.call([resources_path, fingerprint], shell=False)
		except Exception as e:
			print(e)
			raise Exception("Unable to execute the script performing key signing.\n" + resources_path)

		return True

	return False

def delete_file(file_path):
	print('Would you like to delete the original unencrypted message?\n')

	reply = None
	reply = input('Type \'y\' to confirm deletion of the message? : ')
	print()
	if(str(reply) != 'y' and str(reply) != 'Y'):
		print('The file located at %s will NOT be deleted.\n' % file_path)
		return False

	# the utility srm is used to securely remove the unencrypted message.
	try:
		print('Deleting file located at %s\n' % file_path)
		subprocess.check_call([os.getcwd() + "/res/srm.sh", file_path], shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
	except subprocess.CalledProcessError as e:
		print('Unable to delete file using Secure Remove (SRM).')		
		raise Exception('Return code: ' + e.returnCode + '\nErrors returned:\n' + e.stderr + '\n')

	return True

#Functions called exclusively in main

#optional arguments to allow for future use of command line arguments with GPG Package Handler
#def prepare_message(userid=None, keytype=None, keylength=None, subkeytype=None, subkeylength=None, expirydate=None, keyserver=None):
def prepare_message():
	
	RECIPIENT = 0	
	reply = None
	def_keypair = None
	file_found = False
	key_signed = False
	key_server = default_key_server
	recipients = []
	fingerprints = []
	cur_sender_pubkey_path = ""

	print('- Preparing message for encryption.\n')

	#Checking to see if a file is present within the to_encrpyt folder.
	#This check is performed first to avoid completing unnecssary steps.
	for file in os.listdir(encrypt_path):
		if(file.startswith('.') == False):
			file_found = True
	
	if(file_found == False):
		print('- No content to encrypt.\n- Please add content within\n%s\n' % encrypt_path)
		return		

	def_keypair=get_default_keypair()

	if(def_keypair == None):
		print('Unable to retrieve the sender\'s keypair.')
		return		

	prompt = '- Preparing GPG key used to identify sender.\n'
	options = ['1. Use existing default keypair created by ' +  def_keypair.uid, '\n2. Reassign default keypair from existing keys.\n', '3. Create new keypair\n']
	choice = create_menu(prompt,options)

	if(choice == 1):
		print('- Existing key used. Messages will be signed with default key.\n')
		
	if(choice == 2):
		if(reassign_default_keypair() != 0):
			print("Error occured while reassigning default keypair.")
			return
	
		def_keypair=get_default_keypair()
			
	if(choice == 3):
		
		try:
			def_keypair=create_default_keypair(key_server)
		except Exception as e:
			print(e)
			return

	#if(len(gpg.list_keys(secret=True) > 1):
	#	options = ['1. Use existing default keypair created by ' +  def_keypair.uid, '\n2. Reassign default keypair from existing keys.\n', '3. Create new keypair\n']

	#	choice = create_menu(prompt,options)

	#	if(choice == 1):
	#		print('- Existing key used. Messages will be signed with default key.')
		
	#	if(choice == 2):
	#		if(reassign_default_keypair() != 0):
	#			print("Error occured while reassigning default keypair.")
	#			return
	
	#		def_keypair=get_default_keypair()
			
	#	if(choice == 3):
	#		try:
	#			def_keypair=create_default_keypair(key_server)
	#		except Exception as e:
	#			print(e)
	#			return
	#else:
	#	options = ['1. Use existing default keypair created by ' +  def_keypair.uid, '\n2. Create new keypair\n']
	
	#	choice = create_menu(prompt,options)

	#	if(choice == 1):
	#		print('- Existing key used. Messages will be signed with default key.')
		
	#	if(choice == 2):
	#		try:
	#			def_keypair=create_default_keypair(key_server)
	#		except Exception as e:
	#			print(e)
	#			return
				
	print("Info related to sender's key pair.")
	def_keypair.print_info()

	#Collecting recipient information from the sender.
	print('- Determining recipient(s) for message by User ID.\nPlease ensure the receiver has previously created a keypair.\n')
	while(True):
		recipients.append(collect_recipient_info())
		reply = None
		reply = input('Type \'y\' to add another recipient : ')
		print()
		if(str(reply) != 'y' and str(reply) != 'Y'):
			break

	recipient_info = search_for_key(recipients,key_server)

	if(len(recipient_info) == 0):
		print('- Unable to retrieve info on recipient(s).')
		return

	# retrieving recipient(s) public keys from keyserver based on fingerprint of each and import keys into key ring
	for key_info in recipient_info:
		key_signed = False
		try:
			print("- Importing info for %s from %s.\n" % (key_info.uid, key_server))
			key = gpg.recv_keys(key_server, key_info.fingerprint)
			
			if(key.__dict__['results'][RECIPIENT]['fingerprint'] != key_info.fingerprint):
				print("The imported key information does not match the information from the initial search.\n")
				#Deleting the undesired public key. 
				#Deletion will not work if there is a private key existing within the local GPG keyring.			
				gpg.delete_keys(key.__dict__['results'][RECIPIENT]['fingerprint'])   			
				continue

			if(retrieve_valid_key(key_info.fingerprint, verbose=True) == None):			
				print("Unable to send messages to recipient using this key.\n");
				continue
		except Exception as e:
			print(e)
			print('- Unable to import receivers key into local GPG keyring.')
			return			

		#Checking if the receiver's public key has previously been signed by the sender. 
		#Will be useful if an offline mode is developed.
		
		#key_signed = check_for_key_signature(key_info.fingerprint,def_keypair.key_id)		
		
		#Signing recipient's public key the sender's private key.				
		#print()
		#print('- Signing recipient\'s public key to establish trust for future communication.')
		#if(key_signed == False):
		#	try:
		#		if(sign_key(key_info.fingerprint) == True):
		#			print('Trust established as recipient\'s public key has been signed by ' + def_keypair.uid + '.')
		#		else:
		#			print('Trust NOT established with recipient\'s public key.')
					
		#	except Exception as e:
		#		print(e)
		#		print_exit('- Unable to sign recipient\'s public key.')
		#else:
		#	print('Recipient\'s public key has been previously signed by ' + def_keypair.uid + '.') 

		try:
			recip_public_key = gpg.export_keys(key_info.fingerprint, armor=True)
			recipient_pubkey_path = sender_path + 'recipient_public_keys/' + key_info.uid.replace(" ","") + '.asc'
			output_to_file(recipient_pubkey_path, recip_public_key)
			print('- Successfully output public key of %s to file at:\n%s\n' % (key_info.uid,recipient_pubkey_path))
		except Exception as e:
			print(e)
			print('- Unable to output public key of ' + key_info.uid + ' to file.\nCannot sign file to establish trust with recipient.')
			return
		
		# collecting the fingerprints of each recipient to use during encryption
		# the recipients public key is used to encrypt the document, their private key can then be used to decrypt the document
		fingerprints.append(key_info.fingerprint)

	#Export senders public key to file. 
    #Will be useful if an offline mode is developed
	#print()
	#print('- Exporting sender\'s public key to file.\nRecipients can use your public key to establish trust with you as sender.')
	#print('Private Key created, messages can be signed using your private key.')

	#try:
	#	cur_sender_pubkey_path = sender_pubkey_path + '/' + keypair.uids.replace(" ","") +'.asc'
	#	output_to_file(cur_sender_pubkey_path, keypair.public_key)
	#	print()
	#	print('- Sender\'s public key successfully output to\n' + cur_sender_pubkey_path)
	#except Exception as e:
	#	print(e)
	#	print_exit("- Cannot export public key to file.\n- Unable to provide recipient with copy of your public key")
	#

	#extract all fingerprints from list of keyInfo objects
	if(len(fingerprints) == 0):
		print('- Unable to encrypt file. Reciepient info not retrieved.')
		return

	#read name of files folder of files to encrypt
	for file in os.listdir(encrypt_path):
		if(file.startswith('.') == False):
			try:			
				encrypt_file(encrypt_path + file, encrypted_path + file, fingerprints, def_keypair.fingerprint)
				delete_file(encrypt_path + file)
			except Exception as e:
				print(e)
				return

	print('- Encryption process complete.\nSee ' + encrypted_path + ' to find the content to provide to the receiver.\n')

def receive_message():
	
	SENDER = 0
	file_found = False
	key_signed = False	
	def_keypair = None
	sender_key = None
	sender_key_info = None
	sender_name = None
	sender_email = None
	key_server = default_key_server

	print('- Preparing message for decryption.\n')

	#Checking to see if a file is present within the to_encrpyt folder.
	#This check is performed first to avoid completing unnecssary steps.
	for file in os.listdir(encrypted_path):
		if(file.startswith('.') == False):
			file_found = True
	
	if(file_found == False):
		print('- No content to decrypt.\n- Please add content within\n%s\n' % encrypted_path)
		return

	def_keypair=get_default_keypair()

	if(def_keypair == None):
		print('Unable to retrieve the receiver\'s keypair')		
		return

	prompt = '- Preparing GPG key used to identify receiver.\n'
	options = ['1. Use existing default keypair created by ' + def_keypair.uid,'\n2. Reassign default keypair from existing keys.\n']
	choice = create_menu(prompt,options)

	if(choice == 1):
		print('- Existing key used. Messages will be signed with default key.\n')
		
	if(choice == 2):
		if(reassign_default_keypair() != 0):
			print("Error occured while reassigning default keypair.")
			return

		def_keypair=get_default_keypair()
						
	print("Info related to receivers's key pair.")
	def_keypair.print_info()

	print('- Determining sender of the message by User ID.')
	sender_name = input('What is the name of the sender? : ')
	sender_email = input('What is the email of the sender? : ')
	print()

    #query's keyserver and return KeyInfo object if uid is found.
	sender_key_info = search_for_key(format_uid(sender_name,sender_email), key_server)

	if(len(sender_key_info) == 0):
		print('- Unable to retrieve sender\'s key from %s.\n' % key_server)
		return
	
	sender_key_info = sender_key_info[SENDER]
	
	try:
		print("- Importing info for %s from %s.\n" % (sender_key_info.uid, key_server))
		sender_key = gpg.recv_keys(key_server, sender_key_info.fingerprint)
		
		if(sender_key.__dict__['results'][SENDER]['fingerprint'] != sender_key_info.fingerprint):
			print("The imported key information does not match the information from the initial search.\n")
			#Deleting the undesired public key. 
			#Deletion will not work if there is a private key existing within the local GPG keyring.				
			gpg.delete_keys(sender_key.__dict__['results'][SENDER]['fingerprint'])   			
			return

		if(retrieve_valid_key(sender_key_info.fingerprint, verbose=True) == None):			
			print("Unable to receive messages from sender using this key.");
			return

	except Exception as e:
		print(e)
		print('- Unable to import sender\'s key into local GPG keyring.')		
		return

	try:
		print('- Signing sender\'s public key to establish trust.\n')

		#Checking if the sender's public key has previously been signed by the receiver. 
		key_signed = check_for_key_signature(sender_key_info.fingerprint,def_keypair.key_id)	

		if(key_signed == False):
			if(sign_key(sender_key_info.fingerprint) == True):
				
				print('- Trust established!\n%s\'s public key has been signed by %s.\n' % (sender_key_info.uid, def_keypair.uid))
				print('- Sending public key of %s to %s to help establish web of trust\n.' % (sender_key_info.uid, key_server))
				#sending signed public key of sender back to keyserver. 
				send_result=gpg.send_keys(key_server,sender_key_info.fingerprint)		
				#ensuring the key was sent to the key server.	
				errors = str(send_result.__dict__['stderr'])		
				if(errors.find("keyserver send failed") != -1):	
					print('- Error sending public key to keysever.')
			else:
				print('- Trust NOT established with sender\'s public key.\nDecryption process will now stop.\n')
				#Deleting the undesired public key. 
				#Deletion will not work if there is a private key existing within the local GPG keyring.	
				gpg.delete_keys(sender_key_info.fingerprint)
				return
		else:
			print('Sender\'s public key has been previously signed by %s.\n' % def_keypair.uid) 
	except Exception as e:
		print(e)
		print('- Unable to sign sender\'s public key.\nDecryption process will now stop.')
		return

	print('- Outputting the public key of %s to file.\n' % sender_key_info.uid)

	try:
		sender_public_key = gpg.export_keys(sender_key_info.fingerprint, armor=True)
		sender_pubkey_file_path = sender_pubkey_path + sender_key_info.uid.replace(" ","") + '.asc'
		output_to_file(sender_pubkey_file_path, sender_public_key)
		print('Sender\'s public key saved at %s\n' % sender_pubkey_file_path) 

	except Exception as e:
		print(e)
		print('- Unable to output public key of %s to file.\nCannot sign file to establish trust with recipient.' % sender_key_info.uid)
		return

	try:
		print('- Decrypting file(s) received from sender.\n')

		for file in os.listdir(encrypted_path):
			if(file.startswith('.') == False):
				encrypted_file = open(encrypted_path + file, 'rb')
				file = file[:-4]
				result = gpg.decrypt_file(encrypted_file, output=decrypted_path + file)

				#TRUST_LEVELS = {
				#"TRUST_UNDEFINED" : TRUST_UNDEFINED,
				#"TRUST_NEVER" : TRUST_NEVER,
				#"TRUST_MARGINAL" : TRUST_MARGINAL,
				#"TRUST_FULLY" : TRUST_FULLY,
				#"TRUST_ULTIMATE" : TRUST_ULTIMATE,
     		 	#
				# See trust levels above.
     			#}

				#Verifying trust in encrypted file
				if result.trust_level is not None and result.trust_level >= result.TRUST_FULLY:
					print( '- File successfully decrypted.\n%s saved at:\n%s' % (file,decrypted_path))
					print()
				else:
					print('Decryption of ' + file + ' unsuccessful.')
					return					

		#end of encrypted file loop
	except Exception as e:
		print(e)
		print('- Unable to decrpyt message from sender')
		return

#Allows for creation of menus without creating chained if statements.
#prompt - String -  The question asked to the user.
#options - List of String -  The available options to the user.
def create_menu(prompt, options):	
	
	#Building the regular expression used check for valid input within collect_input()
	pattern = '[1-' + str(len(options)) + ']'

	#Appending all options to a single string that is provided to collect_input()
	input_prompt = prompt
	for option in options:
		input_prompt += option
	
	#Adding the line that the user will enter its choice on.
	input_prompt += ' - '
	
	#Presenting the menu until valid input is provided. 
	while(True):
		reply = None
		reply = collect_input(prompt=input_prompt,pattern=pattern)
		if(reply != None):
			break;
	print()
	return int(reply)

def main_menu():
	print("--- MAIN MENU ---\n")	
	
	def_keypair=get_default_keypair()

	if(def_keypair != None):
		print("The key currently in use is:")
		def_keypair.print_info()
	elif( len(gpg.list_keys(secret=True)) == 0):
		print("Please create a keypair.\nA keypair is needed to send or receive packages.\n")
		
		try:		
			create_default_keypair()
		except Exception as e:
			print(e)
			return	
	else:
		prompt = 'What would you like to do?\n'
		options = [ '1. Assign default keypair from existing keys.\n','2. Create and assign new default keypair.\n','3. Exit Program.\n' ]
		functionNames = [ reassign_default_keypair, create_default_keypair , print_exit ]
		choice = create_menu(prompt, options)
		functionToCall = functionNames[choice-1]
		
		try:		
			functionToCall()
		except Exception as e:
			print(e)
			return		

	prompt = 'What would you like to do?\n'
	options = ['1. Send or recieve a package.\n','2. Manage your default keypair.\n','3. Exit program.\n']
	functionNames = [ manage_package_menu, manage_key_menu, print_exit ]
	choice = create_menu(prompt,options)	
	functionToCall = functionNames[choice-1]
	functionToCall()	

def manage_package_menu():
	prompt = 'Are you sending or receiving a package?\n'
	options = [ '1. Sending package.\n','2. Receiving package.\n','3. Return to main menu.\n']
	functionNames = [ prepare_message, receive_message, main_menu ] 
	choice = create_menu(prompt, options)
	functionToCall = functionNames[choice-1]
	functionToCall()

def manage_key_menu():
	prompt = 'What would you like to do?\n'
	options = [ '1. Create and assign new default keypair.\n','2. Reassign default keypair from existing keys.\n','3. Return to main menu.\n' ]
	functionNames = [ create_default_keypair , reassign_default_keypair, main_menu ]
	choice = create_menu(prompt, options)
	functionToCall = functionNames[choice-1]
	
	try:
		functionToCall()
	except Exception as e:
		print(e)
		return	

######################################

#  MAIN

######################################

sys.path.insert(0, os.getcwd() + '/res')

parser = OptionParser()
reply = 1
#NAME and EMAIL are used when parsing user_id command line argument
#NAME = 0
#EMAIL = 1


#reading argument passed in from command line
#home directory, userid (name,email), keytype, keylength, subkeytype, subkeylength, expiry date, recipients (recip1, recip2)
parser.add_option('-l', dest='home_directory', help='-p Path of current users home directory leading to where GPG is installed. run.sh provides this argument automatically')
#parser.add_option('-u', dest='user_id', help='-u Name,Email')
#parser.add_option('-e', dest='expiry_date',help='-e Expiry Date. Format: yyyy-mm-dd')
#parser.add_option('-r', dest='recipients',help='r RecipientList Format: recip1, recip2')
#parser.add_option('--kt', dest='key_type', help='--ky Key Type')
#parser.add_option('--kl', dest='key_length', help='--kl Key Length. Max = : :')
#parser.add_option('--skt', dest='subkey_type', help='skt Subkey Type')
#parser.add_option('--skl', dest='subkey_length', help='--skl Subkey Length. Max = : :')

(options, args) = parser.parse_args()

#if(options.user_id != None):
#	#Vinay Sajip <vinay_sajip@hotmail.com>'
#	id_parts = user_id.split(',')
#	user_id = format_uid(id_parts[NAME], id_parts[EMAIL])

#if(options.recipients!= None):
#	recipient_list  = recipients.split(',')

#should be able to avoid retrieving and passing the home directory from the installation script
#gpg_path = '~/.gnupg'
gpg_path = options.home_directory + '/.gnupg'
#print(gpg_path)

gpg = gnupg.GPG(gnupghome=gpg_path)
gpg.encoding = 'utf-8'

while(True):
	main_menu()
