#!/usr/bin/env python3
import socket
from threading import Thread
from Crypto.PublicKey import RSA
from datetime import datetime
from src.socks import receive, receive_session, receive_aes, send, PORT
from src.keys import login
from src.contacts import load_contacts, save_contacts
from src.ui import *


def print_banner():
    """Prints the entry banner."""
    print_green("/////////////////////")
    print_green("// s l y t h e r ////")
    print_green("////// s e r v e r //")
    print_green("/////////////////////")


def get_contact_name(ip, contacts):
    """
    Given an IP address, finds the corresponding contact name.

    Args:
        ip: The ip address to match with a contact name.
        contacts: The contacts dictionary to search.

    Return:
        The name of the contact if the IP is known, otherwise the IP.
    """
    for contact in contacts:
        if contacts[contact]["ip"] == ip and contact != ip:
            return contact
    return ip


def handle_client(sock, addr, public, private):
    """
    Thread that receives a message from a connection.

    Args:
        sock: The socket the client has connected on.
        addr: Tuple of the IP address and port of the connection.
        public: The public key of this user.
        private: The private key of this user.
    """
    print("Handling client")
    contacts = load_contacts()
    display_name = get_contact_name(addr[0], contacts)
    print_green("New connection from {}!".format(display_name))
    
    print(" > Performing key exchange...")
    client_public = RSA.import_key(receive(sock))
    print("    :  Received public key.")
    send(sock, public.export_key())
    print("    :  Sent public key.")

    print(" > Receiving message...")
    try:
        session_key = receive_session(sock, client_public, private)
        print("    : Received session key.")
        message = receive_aes(sock, client_public, session_key)
        print("    : Received message.")
    except ValueError as e:
        print_red("Error receiving message.")
        print(e)
    else:
        print(" > Storing message...")
        message_receipt = { "time": datetime.now().strftime("%m/%d/%y %I:%M%p"), 
                            "to": "me", 
                            "from": display_name, 
                            "contents": message.decode() }
        if display_name in contacts:
            contacts[display_name]["messages"].append(message_receipt)
        else:
            contacts[display_name] = {  "ip": addr[0],
                                        "fingerprint": None,    # Tweak when fingerprint time
                                        "messages": [message_receipt]}
        save_contacts(contacts)
    finally:
        print(" > Closing connection...")
        sock.close()
        

if __name__ == "__main__":
    print_banner()
    # Load keys
    public, private = login()


    # Bind socket
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
            server_addr = ("0.0.0.0", PORT)
            sock.bind(server_addr)
            sock.listen(5)

            # Acceptance loop
            print("Listening for connections...")
            while True:
                try:
                    connection, addr = sock.accept()
                    print("Connection")
                    connection_thread = Thread( target=handle_client, 
                                                args=(connection, addr, public, private))
                    connection_thread.start()
                except KeyboardInterrupt:
                    if confirm("\nAre you sure you'd like to close slyther-server? (Y/n) "):
                       break 
    except OSError:
        print_red("Error: Failed to start slyther-server: Port {} in use.".format(PORT))


