import sys
import os
import socket
import SocketServer
import threading
import RSA

def main(): # Checks for correct input arguments and runs the appropriate methods
    if (len(sys.argv) < 3):
        print 'Usage: python serverclient.py <server|client> <port>\n'
        return -1
    else:
        if sys.argv[1].lower() == 'server':
            Server(sys.argv[2])
        elif sys.argv[1].lower() == 'client':
            Client(sys.argv[2])
        else:
            print 'Unrecognized argument: ', sys.argv[1]
            return -1
    return 0

def Server(port): # Creates the server instance, sets it up
    host = 'localhost'
    port = int(port)
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind((host, port))
    server.listen(5)
    # blocking call to accept()
    print 'Waiting for partner to join conversation...\n'
    (conn, client_addr) = server.accept()
    print 'Client connected: ', client_addr[0]
    # wait to receive client's public key
    key = conn.recv(1024)
    key = key.split(',')
    keyTuple = (key[0], key[1])
    print 'Client\'s Public Key received'
    e, d, c = RSA.keygen()
    sendPublic = str(d) + ',' + str(c)
    conn.send(sendPublic)
    print 'Public Key sent'
    privateTuple = (e, c)
    print 'Type your message below and hit enter to send. Type \'EXIT\' to end conversation.\n'
    ReadThread = Thread_Manager('read', conn, keyTuple, None)
    WriteThread = Thread_Manager('write', conn, None, privateTuple)
    ReadThread.start()
    WriteThread.start()
    # wait until client dc's
    ReadThread.join()
    print 'Your partner has left the conversation. Press any key to continue...\n'
    # stop the write thread
    WriteThread.stopWrite()
    WriteThread.join()
    # shut down client connection
    try:
        conn.shutdown(socket.SHUT_RDWR)
        conn.close()
    except:
        # connection already closed
        pass
    # shut down server
    print 'Shutting server down...'
    server.shutdown(socket.SHUT_RDWR)
    server.close()
    return 0

def Client(port): # Creates the client instance, sets up the client
    host = 'localhost'
    port = int(port)
    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client.connect((host, port))
    e, d, c = RSA.keygen()
    sendPublic = str(d) + "," + str(c)
    client.send(sendPublic)
    print 'Public key sent'
    key = client.recv(1024)
    key = key.split(',')
    keyTuple = (key[0], key[1])
    print 'Server\'s Public Key received'
    privateTuple = (e, c)
    print 'Connected to chat. Type your messages below and hit enter to send. Send \'EXIT\' to leave the conversation.\n'
    ReadThread = Thread_Manager('read', client, keyTuple, None)
    WriteThread = Thread_Manager('write', client, None, privateTuple)
    ReadThread.start()
    WriteThread.start()
    ReadThread.join()
    print 'Your partner has left the conversation. Press any key to continue...\n'
    # stop the write thread
    WriteThread.stopWrite()
    WriteThread.join()
    # shut down client connection
    try:
        client.shutdown(socket.SHUT_RDWR)
        client.close()
    except:
        # connection already killed
        pass

class Thread_Manager(threading.Thread): # Creates threads for asynchronoues reading and writing
    def __init__(self, action, conn, public, private): # Constructor for Thread_Manager
        threading.Thread.__init__(self)
        self.action = action.lower()
        self.conn = conn
        self.dowrite = True
        self.exitcode = 'EXIT'

        if public is not None:
            self.setPublic(public)
        if private is not None:
            self.setPrivate(private)

    def run(self): # Invoked when new thread is executed
        if (self.action == 'read'):
            self.read()
        else:
            self.write()

    def setPublic(self, public): # Sets public key from other party for decryption
        self.public = public

    def setPrivate(self, private): # Sets private key for encryption
        self.private = private

    def stopWrite(self): # Terminates the write loop
        self.dowrite = False

    def decrypt(self, buff): # Decrypts input integer list into sentences
        words = buff.split(",")
        decrypted_data = ""
        for i in range(0, len(words) - 1):
            decrypted_data += str(RSA.decode(RSA.endecrypt(words[i], self.public[0], self.public[1])))
        return decrypted_data

    def read(self): # Responsible for reading in data from the client and displaying stdout
        buff = self.conn.recv(1024)
        buff = self.decrypt(buff)
        while buff.strip() != self.exitcode and len(buff) > 0:
            print 'Message received: ', buff.strip()
            buff = self.conn.recv(1024)
            buff = self.decrypt(buff)
        # client disconnected
        self.stopWrite

    def encrypt(self, data):
        encrypted_data = ""
        for i in range(0, len(data)):
            encrypted_data += str(RSA.endecrypt(ord(data[i]), self.private[0], self.private[1])) + ","
        return encrypted_data

    def write(self): # Responsible for reading in data from stdin and sending to client
        while self.dowrite:
            data = sys.stdin.readline()
            data = self.encrypt(data)
            self.conn.send(data)
            if (data.strip() == self.exitcode):
                self.conn.shutdown(socket.SHUT_RDWR)
                self.conn.close()
                self.dowrite = False

# Entry point
if __name__ == "__main__":
    sys.exit(main())