import socket 
from threading import Thread 
from socketserver import ThreadingMixIn 
from utils.keys_manager_1 import *  
from utils.symmetric_keys_manager_1 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):
    got_pub_key_client = False  
    able_to_retrieve_session_id = False 
    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): 
        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-----" 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 [ ~ ]") 
                        session_id = msg
                        (symmetric_key,client_pub_key) = 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 ]") 
                            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... [ :) ]")  
                            # --- 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 = gen_sym_key_and_save(public_key_client) 
                    print("[DEBUG] No existing Session, will now initiate new Session [ OK ]") 
                    print("[DEBUG] I'm all yours... [ :) ]")  
                    
            print("Server received data:", data) 
            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()