#===========================================================================================
#  _____ _____ _   _    _    __  __  ___  ____ _____ _____ ____ _   _ 
# |_   _| ____| \ | |  / \  |  \/  |/ _ \|  _ \_   _| ____/ ___| | | |
#   | | |  _| |  \| | / _ \ | |\/| | | | | |_) || | |  _|| |   | |_| |
#   | | | |___| |\  |/ ___ \| |  | | |_| |  _ < | | | |__| |___|  _  |
#   |_| |_____|_| \_/_/   \_\_|  |_|\___/|_| \_\|_| |_____\____|_| |_|
#
#	RSA Keys Manager V1 (Python 3.6)
#
#   Sources: 
#    https://nitratine.net/blog/post/asymmetric-encryption-and-decryption-in-python/
#    https://stackoverflow.com/questions/8933237/how-to-find-if-directory-exists-in-python 
#    https://stackoverflow.com/questions/273192/how-can-i-safely-create-a-nested-directory 
#
#===========================================================================================

# Basic cryptography tools 
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.asymmetric import rsa

# To format created key obj into text to be used/exported 
from cryptography.hazmat.primitives import serialization 

# To get user home directory 
from pathlib import Path


def generate_private_key(): 
	# Generate Private key
	private_key = rsa.generate_private_key(
		public_exponent=65537,
		key_size=2048,
		backend=default_backend()
	) 
	save_key(private_key, 0) 
	return private_key 


def generate_public_key(private_key): 
	# Get Public key from Private 
	public_key = private_key.public_key() 
	save_key(public_key, 1) 
	return public_key 


def private_key_serializer(private_key): 
	# Key serialization 
	private_key_pem = private_key.private_bytes(
		encoding=serialization.Encoding.PEM,
		format=serialization.PrivateFormat.PKCS8,
		encryption_algorithm=serialization.NoEncryption()
	)
	return private_key_pem 


def public_key_serializer(public_key): 
	# Key serialization 
	public_key_pem = public_key.public_bytes(
		encoding=serialization.Encoding.PEM,
		format=serialization.PublicFormat.SubjectPublicKeyInfo
	)
	return public_key_pem 


def save_key(key, key_type):
	keys_dir_path = Path.home() / '.auth_server_test' 
	# Path exists ? If not we will create it 
	keys_dir_path.mkdir(exist_ok=True)
	if(key_type == 0): 
		key_pem = private_key_serializer(key)
		key_path = keys_dir_path / 'id_rsa' 
	elif(key_type == 1): 
		key_pem = public_key_serializer(key) 
		key_path = keys_dir_path / 'id_rsa.pub' 
	else: 
		print("ERROR: Key type error") 
		exit()

	with open(str(key_path), 'wb') as f:
		f.write(key_pem)
	 

def generate_keys():
	private_key = generate_private_key() 
	public_key  = generate_public_key(private_key) 
	return (private_key,public_key)


def check_for_existing_keys(): 
	keys_dir_path = Path.home() / '.auth_server_test' 
	private_key_path = keys_dir_path / 'id_rsa' 
	public_key_path = keys_dir_path / 'id_rsa.pub' 
		
	# =========================================================
	# Private key file exists ? Load it
	# --- Public key file exists ? Load it
	# --- No Public key file ? Generate it from the Private
	# No Private key file ? Generate both keys (save as well) 
	# =========================================================

	if(private_key_path.exists()):
		print('Find private key from file [ OK ]') 
		# Load private key from file
		print('Load private key from file [ ... ]') 
		with open(str(private_key_path), "rb") as key_file:
			private_key = serialization.load_pem_private_key(
				key_file.read(),
				password=None,
				backend=default_backend()
			)
		print('Load private key from file [ OK ]') 
		if(public_key_path.exists()): 
			print('Find public key from file [ OK ]') 
			# Load public key from file 
			print('Load public key from file [ ... ]') 
			with open(str(public_key_path), "rb") as key_file:
				public_key = serialization.load_pem_public_key(
					key_file.read(),
					backend=default_backend()
				)
			print('Load public key from file [ OK ]') 
			# Return keys 
			return (private_key,public_key) 
		else:
			print('Find public key from file [ FAIL ]') 
			print('Generating public key from private key [ ... ]') 
			public_key = generate_public_key(private_key) 
			print('Generating public key from private key [ OK ]') 
			return (private_key,public_key) 
	else: 
		print('Find private key from file [ FAIL ]')
		print('New keys will be generated') 
		return generate_keys()  


check_for_existing_keys()