slyther-server

#!/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 *
from src.fingerprints import create_fingerprint, verify_fingerprint


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_id(ip, contacts):
    """
    Given an IP address, finds the corresponding contact ID.

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

    Return:
        The ID of the contact if the IP is known, otherwise the IP.
    """
    for contact_id in contacts:
        if contacts[contact_id]["ip"] == ip:
            return contact_id
    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.
    """
    contacts = load_contacts(private)
    contact_id = get_contact_id(addr[0], contacts)
    
    if contact_id in contacts:
        print_green("New connection from {}!".format(contacts[contact_id]["name"]))
    else:
        print_green("New connection from {}!".format(contact_id))
    
    try:
        print(" > Performing key exchange...")
        client_public = RSA.import_key(receive(sock))
        print("    : Received public key.")
        
        # Check fingerprint
        if contact_id in contacts and contacts[contact_id]["fingerprint"] is not None:
            print(contacts[contact_id]["fingerprint"])
            if not verify_fingerprint(client_public, contacts[contact_id]["fingerprint"]):
                print_yellow("    : Fingerprint mismatch. Untrusted.")

        send(sock, public.export_key())
        print("    : Sent public key.")

        print(" > Receiving message...")
        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)
    except OSError:
        print_red("    : Connection lost. Message not recieved.")
    else:
        print(" > Storing message...")

        message_receipt = { "time": datetime.now().strftime("%m/%d/%y %I:%M%p"), 
                            "recieved": True,
                            "contents": message.decode() }
        if contact_id in contacts:
            # Edit message if untrusted 
            if contacts[contact_id]["fingerprint"] is not None and not verify_fingerprint(client_public, contacts[contact_id]["fingerprint"]):
                message_receipt = { "time": datetime.now().strftime("%m/%d/%y %I:%M%p"), 
                                    "recieved": True,
                                    "contents": "!UNTRUSTED! {} !UNTRUSTED!".format(message.decode()) }
            contacts[contact_id]["messages"].append(message_receipt)
        else:
            contacts[contact_id] = {  "name": contact_id,
                                        "ip": addr[0],
                                        "fingerprint": create_fingerprint(client_public),
                                        "messages": [message_receipt]}

        save_contacts(contacts, private)
    finally:
        print(" > Closing connection...\n")
        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()
                    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))