# Symmetric crypto lib
from cryptography.fernet import Fernet

#UniqueID 
import uuid 

# To get user home directory 
from pathlib import Path
import os 
# Home made RSA Keys lib 
from utils.rsa_tenamortech_utils import *


def gen_rand_session_id(): 
    return str(uuid.uuid4()) 

# --- Used by Server 
def gen_sym_key_and_save(public_key_client): 
    key = Fernet.generate_key() 

    root_file_path = Path.home() / '.auth_server_test/server'
    path_public_keys_clients = root_file_path / 'public_keys_clients'      
    # Path exists ? If not we will create it 
    path_public_keys_clients.mkdir(exist_ok=True)

    session_id = gen_rand_session_id() 
 
    file = open(str(root_file_path) + '/.known_clients', 'ab') # changed from 'wb' to 'w+' ... 
    file.write((session_id + " ").encode('utf-8'))  
    file.write(key + (" ").encode('utf-8')) # The key is type bytes still
    file.write(('0\n').encode('utf-8')) 
    file.close() 

    # --- Store the client public Key for next time.
    # --- Public keys are stored in '.auth_server_test/server/pubilic_keys_clients/session_id.pub'
    public_key_client_pem = public_key_serializer(public_key_client) 
    key_path = path_public_keys_clients / (str(session_id) + '.pub') 

    with open(str(key_path), 'wb') as f:
        f.write(public_key_client_pem) 
    
    return (key,session_id)  

# --- Used by server
def reload_session_sym_key(sessionid):
    root_file_path = Path.home() / '.auth_server_test/server'
    path_public_keys_clients = root_file_path / 'public_keys_clients'      
    # Path exists ? If not we will create it 
    path_public_keys_clients.mkdir(exist_ok=True)
    
    session_id_found = False 
    
    # --- Open session file
    with open(str(root_file_path) + "/.known_clients") as f:
        line = f.readline().rstrip() 
        while(line and session_id_found == False): 
            # each line contains (session_id, sym_key, sym_key_nb_use) 
            # Those infos are spaced with "space" char, so we use 'split' fuction to get an array 
            cur_client_infos = line.split() 
            # If we find the given sessionID: 
            #print(cur_client_infos[0] + "|" + sessionid.decode('utf-8'))
            if(cur_client_infos[0] == sessionid.decode('utf-8')): # cur_client_info[0] is the session_id stored in the cur line  
                session_id_found = True  
            line = f.readline().rstrip() 
    
    # --- We found the given sessionID, so we will now load the corresponding client public key we previously stored.
    if session_id_found == True: 
        with open(str(path_public_keys_clients) + "/" + cur_client_infos[0] + '.pub', "rb") as key_file:
            public_key_client = serialization.load_pem_public_key(
                key_file.read(),
                backend=default_backend()
            )
        # --- Return the Symmetric Key used with this client and the PEM formated client public Key 
        return (cur_client_infos[1],public_key_client) 
    else: 
        # Return Error tuple 
        return (-1,-1) 


# --- Used by client
def reload_session_by_host(host): 
    root_file_path = Path.home() / '.auth_server_test/client'
    path_public_keys_servers = root_file_path / 'public_keys_servers'      
    # Path exists ? If not we will create it 
    path_public_keys_servers.mkdir(exist_ok=True)
    
    host_found = False 
    
    # --- Open session file 
    if os.path.isfile(str(root_file_path) + "/.known_hosts"):    
        with open(str(root_file_path) + "/.known_hosts") as f:
            line = f.readline().rstrip() 
            while(line and host_found == False): 
                # each line contains (host session_id, sym_key) 
                # Those infos are spaced with "space" char, so we use 'split' fuction to get an array 
                cur_server_infos = line.split() 
                # If we find the given host: 
                if(cur_server_infos[0] == host): # cur_client_info[0] is the host stored in the cur line 
                    host_found = True  
                line = f.readline().rstrip() 
    
    # --- We found the given host, so we will now load the corresponding server public key we previously stored.
    if host_found == True: 
        with open(str(path_public_keys_servers) + "/" + cur_server_infos[1] + '.pub', "rb") as key_file:
            public_key_server = serialization.load_pem_public_key(
                key_file.read(),
                backend=default_backend()
            )
        # --- Return the SessionID, the Symmetric Key and the PEM formated server public Key used with this server
        return (cur_server_infos[1], cur_server_infos[2], public_key_server) 
    else: 
        # Return Error tuple 
        return (-1,-1,-1) 
        

# --- Used by client 
def save_sym_key_by_host(host, session_id, sym_key, public_key_server): 
    root_file_path = Path.home() / '.auth_server_test/client'
    path_public_keys_servers = root_file_path / 'public_keys_servers'      
    # Path exists ? If not we will create it 
    path_public_keys_servers.mkdir(exist_ok=True)
    
    line_to_add = str(host) + " " + session_id + " " + sym_key
    # dont forget to save public key server in separate file !
    
    file = open(str(root_file_path) + '/.known_hosts', 'ab') # changed from 'wb' to 'ab' ... 
    file.write(line_to_add.encode('utf-8'))  
    file.write(('\n').encode('utf-8')) 
    file.close() 

    # --- Store the server public Key for next time.
    # --- Public keys are stored in '.auth_server_test/client/public_keys_servers/session_id.pub'
    public_key_server_pem = public_key_serializer(public_key_server) 
    key_path = path_public_keys_servers / (str(session_id) + '.pub') 

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

# --- Used by Client and Server     
def encrypt_msg_symmetric(msg, key): 
    msg = msg.encode() 
    f_key = Fernet(key)
    enc_msg = f_key.encrypt(msg)

#print(encrypted)

#decrypted = f.decrypt(encrypted) 

#print(decrypted)  

#gen_sym_key_and_save() 
#print(reload_session_sym_key("26c9f89a-f8e1-4b22-83bd-417c5047e527"))