from donna25519 import *
import socket
import threading
import struct

TCP_IP = '127.0.0.1'
MAX_CLIENTS = 10

class ClientRecord:
    """Structure for keeping record of a client's public keys"""
    def __init__(self, I, S, O_queue, clientsocket):
        self.I = I
        self.S = S
        self.O_queue = O_queue
        self.s = clientsocket

class Server:
    def __init__(self, port):
        self.port = port
        self.next_id = 1

        self.registered_clients = {}

    def start(self):
        """Start listening to client connections"""
        print("Starting server on port %d" % self.port)
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.bind((TCP_IP, self.port))
        s.listen(MAX_CLIENTS)
        while True:
            c, addr = s.accept() # connection with new client
            print("New client from address " + str(addr))
            threading.Thread(target = self.on_new_client, args = (c, addr)).start()

        s.close()
        
    def on_new_client(self, clientsocket, addr):
        """Runs when a new client connects to the server"""
        I = PublicKey(clientsocket.recv(32))
        print("Received I = " + I.public.hex())
        S = PublicKey(clientsocket.recv(32))
        print("Received S = " + S.public.hex())
        o_num = struct.unpack('!i', clientsocket.recv(4))[0]
        O_queue = []
        for i in xrange(o_num):
            O_queue.append(PublicKey(clientsocket.recv(32)))
            print("Received O = " + O_queue[i].public.hex())

        user_id = self.register_client(I, S, O_queue, clientsocket)
        clientsocket.send(struct.pack('!i', user_id))

        while True:
            cmd = clientsocket.recv(4)
            if cmd == "GENS": # GENerate Session
                raw_id = clientsocket.recv(4)
                rec_user_id = struct.unpack('!i', raw_id)[0]
                if rec_user_id not in self.registered_clients:
                    clientsocket.send("FAIL")
                    continue
                clientsocket.send(b"GOOD")
                # Send Irec, Srec, Orec to initiator
                Irec = self.registered_clients[rec_user_id].I
                Srec = self.registered_clients[rec_user_id].S
                Orec = self.registered_clients[rec_user_id].O_queue.pop()
                clientsocket.send(Irec.public)
                clientsocket.send(Srec.public)
                clientsocket.send(Orec.public)

                # get Eini from initiator
                Eini = PublicKey(clientsocket.recv(32))

                rec_clientsocket = self.registered_clients[rec_user_id].s
                rec_clientsocket.send(b"NEWS") # NEW Session
                rec_clientsocket.send(struct.pack('!i', user_id))
                # Send Eini, Iini, Orec to recipient
                rec_clientsocket.send(Eini.public)
                rec_clientsocket.send(I.public)
                rec_clientsocket.send(Orec.public)
            if cmd == "SEND": # SEND Message
                # Get recipient used id
                raw_id = clientsocket.recv(4)
                rec_user_id = struct.unpack('!i', raw_id)[0]

                rec_clientsocket = self.registered_clients[rec_user_id].s
                rec_clientsocket.send(b"MESG") # send message tag

                # Send sender user id
                rec_clientsocket.send(struct.pack('!i', user_id))

                # Forward sender ephemereal key
                eph_s = clientsocket.recv(32)
                rec_clientsocket.send(eph_s)

                # Forward message length
                raw_mlen = clientsocket.recv(4)
                rec_clientsocket.send(raw_mlen)

                # Forward message
                mlen = struct.unpack('!i', raw_mlen)[0]
                raw_msg = clientsocket.recv(mlen)
                rec_clientsocket.send(raw_msg)

    def register_client(self, I, S, O_queue, clientsocket):
        """Registers new client's I, S, and O_queue and returns the new client's id"""
        new_client = ClientRecord(I, S, O_queue, clientsocket)

        user_id = self.next_id
        self.next_id += 1

        self.registered_clients[user_id] = new_client

        return user_id