# 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 = input("tcpClientA: Enter message/ Enter exit:").encode('utf-8') 
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:  
            #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] 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()