#===========================================================================================
#  _____ _____ _   _    _    __  __  ___  ____ _____ _____ ____ _   _ 
# |_   _| ____| \ | |  / \  |  \/  |/ _ \|  _ \_   _| ____/ ___| | | |
#   | | |  _| |  \| | / _ \ | |\/| | | | | |_) || | |  _|| |   | |_| |
#   | | | |___| |\  |/ ___ \| |  | | |_| |  _ < | | | |__| |___|  _  |
#   |_| |_____|_| \_/_/   \_\_|  |_|\___/|_| \_\|_| |_____\____|_| |_|
#
#   Minimalist Auth Server V3 (Python 3.6)
#	(This file contains the client code)
#
#   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 
#
#
#	ROMANET Valentin
#===========================================================================================


# Python TCP Client A
import socket 

# 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

# Home made RSA Utils 
from utils.rsa_tenamortech_utils import * 
from utils.client_keys_manager import * 
from utils.symmetric_keys_manager_1 import * 

host = socket.gethostname() 
port = 2004
BUFFER_SIZE = 2000 
MESSAGE = ""  
tcpClientA = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 
tcpClientA.connect((host, port))
got_pub_key_server = False 

(private_key,public_key) = check_for_existing_keys() 
public_key_pem = public_key_serializer(public_key) 


def de_serialize_pub_key(public_key_pem): 
    return serialization.load_pem_public_key(public_key_pem,backend=default_backend())  
 
while MESSAGE != 'exit':
    data = tcpClientA.recv(BUFFER_SIZE) 
    if not got_pub_key_server: 
        print("[DEBUG] Checking for SessionID [ ... ]")
        (session_id,symmetric_key,public_key_server) = reload_session_by_host(host)
        print("[DEBUG] Checking for SessionID [ OK ]")
        if(session_id != -1): 
            tcpClientA.send(encrypt_msg(session_id.encode('utf-8'), public_key_server))
            got_pub_key_server = True  
            print("[DEBUG] SessionID Found ! [ OK ]")
        else:  
            print("[DEBUG] SessionID Found ! [ FAIL ]")
            print("[DEBUG] Waiting for new SessionID [ ... ]")
            #if find sessionID load it and negociate with server
            #else send public key ...  
            print("[DEBUG] Receiving Server Public Key [ ... ]")
            public_key_server = de_serialize_pub_key(data) 
            got_pub_key_server = True  
            print("[DEBUG] Receiving Server Public Key [ OK ]")
            # Now we have the pub key of the server, we will send our pub key too 
            print("[DEBUG] Sending current Client Public Key [ ... ] ")
            tcpClientA.send(public_key_pem)
            print("[DEBUG] Sending current Client Public Key [ OK ] ")
            print("[DEBUG] Waiting for new Symmetric Key and New SessionID [ ... ] ") 
            data = tcpClientA.recv(BUFFER_SIZE) 
            new_symmetric_key_and_session_id = decrypt_msg(data, private_key) 
            new_symmetric_key_and_session_id = new_symmetric_key_and_session_id.decode('utf-8')
            new_symmetric_key_and_session_id = new_symmetric_key_and_session_id.split('|')
            new_symmetric_key = new_symmetric_key_and_session_id[0] 
            new_session_id    = new_symmetric_key_and_session_id[1] 
            
            print("[DEBUG] Waiting for new Symmetric Key and New SessionID [ OK ] ") 
            print("[DEBUG] Symmetric Key: [" + str(new_symmetric_key) + "]") 
            print("[DEBUG] SessionID    : [" + new_session_id + "]") 
            print("[DEBUG] Waiting for new SessionID [ OK ]")
            
            print("[DEBUG] Saving basic HandShake infos [ ... ] ")
            save_sym_key_by_host(host, new_session_id, new_symmetric_key, public_key_server) 
            print("[DEBUG] Saving basic HandShake infos [ OK ] ")
        data = tcpClientA.recv(BUFFER_SIZE)
        decrypted_data = decrypt_msg(data, private_key)
        print(decrypted_data.decode('utf-8')) 
        
    else:  
        print(" Client2 received data:", data) 
        MESSAGE = input("Working !").encode('utf-8') 
        tcpClientA.send(encrypt_msg(MESSAGE, public_key_server))     

tcpClientA.close()