#===========================================================================================
#  _____ _____ _   _    _    __  __  ___  ____ _____ _____ ____ _   _ 
# |_   _| ____| \ | |  / \  |  \/  |/ _ \|  _ \_   _| ____/ ___| | | |
#   | | |  _| |  \| | / _ \ | |\/| | | | | |_) || | |  _|| |   | |_| |
#   | | | |___| |\  |/ ___ \| |  | | |_| |  _ < | | | |__| |___|  _  |
#   |_| |_____|_| \_/_/   \_\_|  |_|\___/|_| \_\|_| |_____\____|_| |_|
#
#    RSA Utils 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

# Encryption 
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding

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


def generate_public_key(host, private_key): 
    # Get Public key from Private 
    public_key = private_key.public_key() 
    save_key(public_key, 1, host) 
    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
    )
    #print(public_key_pem)
    return public_key_pem 


def save_key(key, key_type, host): 
    root_file_path = '.auth_server_test/'
    if host == 0: 
        path = root_file_path + 'server' 
    elif host == 1: 
        path = root_file_path + 'client' 
    keys_dir_path = Path.home() / path  
    # 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(host):
    private_key = generate_private_key(host) 
    public_key  = generate_public_key(host,private_key) 
    return (private_key,public_key)

def encrypt_msg(msg,public_key): 
    enc_msg = public_key.encrypt(
        msg,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )   
    )   
    return enc_msg 


def decrypt_msg(enc_msg,private_key): 
    msg = private_key.decrypt(
        enc_msg,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )
    return msg