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

import socket 
from threading import Thread 
from socketserver import ThreadingMixIn 
#from utils.keys_manager_1 import *  
from utils.server_keys_manager import *  
from utils.symmetric_keys_manager_1 import *
from utils.rsa_tenamortech_utils import *
# --- Init keys --- 


def de_serialize_pub_key(public_key_pem): 
    return serialization.load_pem_public_key(public_key_pem,backend=default_backend())  

# Multithreaded Python server : TCP Server Socket Thread Pool
class ClientThread(Thread):
    def __init__(self,ip,port, private_key, public_key): 
        Thread.__init__(self) 
        self.ip = ip 
        self.port = port
        #print(public_key_serializer(public_key))  
        print("[+] New server socket thread started for " + str(ip) + ":" + str(port))  
        public_key_pem = public_key_serializer(public_key)  
        conn.send(public_key_pem) 
 
    def run(self): 
        got_pub_key_client = False  
        able_to_retrieve_session_id = False 
        while True : 
            data = conn.recv(2048)
            if not got_pub_key_client: # we might receive a pub_key in clear, or an encrtpted session_id
                print("[DEBUG] Waiting for HandShake [ ... ]")
                # --- Text was clear and client sent pub_key 
                if ("-----BEGIN PUBLIC KEY-----").encode('utf-8') in data:
                    print("[DEBUG] Received Public Key from Client [ OK ]") 
                    public_key_client = de_serialize_pub_key(data) 
                    got_pub_key_client = True  
                # --- Decrypt data and try to load session_id etc
                else:
                    print("[DEBUG] No HandShake, got SessionID from Client")
                    print("[DEBUG] Waiting for SessionID validation [ ... ]")    
                    msg = decrypt_msg(data, private_key)
                    if(len(msg) == 36): # SessionID length should be 36 char.
                        print("[DEBUG] SessionID format looks correct [ ~ ]")
                        print("[DEBUG] (Given SessionID: [" + msg.decode('utf-8') + "]")
                        session_id = msg
                        (symmetric_key,public_key_client) = reload_session_sym_key(session_id) 
                        
                        # --- Unable to retrieve sessionID and public Key was not sent
                        # --- Abort connection now  
                        if(symmetric_key == -1): 
                            print("[DEBUG] Waiting for SessionID validation [ FAIL ]") 
                            print("[DEBUG] Invalid SessionID, will now abort negociation [ FAIL ]") 
                            conn.close() 
                            break 
                        else: 
                            able_to_retrieve_session_id = True 
                            print("[DEBUG] Waiting for SessionID validation [ OK ]") 
                            print("[DEBUG] SessionID is valid, symmetric_key and current Client Public key were retrieved successfully")  
                            print("[DEBUG] I'm all yours... [ :) ]")  
                            conn.send(encrypt_msg("login:".encode('utf-8'), public_key_client))  
                            # --- SessionID that the client sent is valid, we got all the required infos
                # --- Client sent did not send sessionID but sent his public Key 
                # --- Create a new session for it, (gen sessionID, create new symmetric Key, store his public key) 
                if able_to_retrieve_session_id == False:
                    print("[DEBUG] No existing Session, will now initiate new Session [ ... ]")  
                    # --- This function will create a sessionID, return a symmetric key and store the pub key
                    (symmetric_key,session_id) = gen_sym_key_and_save(public_key_client)
                    print("[DEBUG] Sending encrypted Symmetric Key to Client [ ... ]") 
                    sym_key_and_session_id = (str((symmetric_key).decode('utf-8')) + "|" + session_id).encode('utf-8')  
                    conn.send(encrypt_msg(sym_key_and_session_id, public_key_client))
                    print("[DEBUG] Sending encrypted Symmetric Key to Client [ OK ]") 
                    print("[DEBUG] I'm all yours... [ :) ]")
                    conn.send(encrypt_msg("login:".encode('utf-8'), public_key_client))  
                    
            MESSAGE = input("Multithreaded Python server : Enter Response from Server/Enter exit:").encode('utf-8') 
            if MESSAGE == 'exit':
                break
            conn.send(MESSAGE)  # echo

# Multithreaded Python server : TCP Server Socket Program Stub
TCP_IP = '0.0.0.0' 
TCP_PORT = 2004 
BUFFER_SIZE = 20  # Usually 1024, but we need quick response 

tcpServer = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 
tcpServer.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 
tcpServer.bind((TCP_IP, TCP_PORT)) 
threads = [] 

# --- Init keys --- 
(private_key, public_key) = check_for_existing_keys() 
#print(public_key_serializer(public_key))  

while True: 
    tcpServer.listen(4) 
    print("Multithreaded Python server : Waiting for connections from TCP clients...")
    (conn, (ip,port)) = tcpServer.accept() 
    newthread = ClientThread(ip,port,private_key,public_key) 
    newthread.start() 
    threads.append(newthread) 
 
for t in threads: 
    t.join()