import sys
import socket
import select
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP
from base64 import b64decode

with open('key_private.pem') as data:
    key_1 = data.read()

with open('key_public.pem') as data:
    key_2 = data.read()

def encrypt_RSA(key, message):
    rsakey = RSA.importKey(key)
    rsakey = PKCS1_OAEP.new(rsakey)
    encrypted = rsakey.encrypt(message)
    return encrypted.encode('base64')

def decrypt_RSA(key, package):
    rsakey = RSA.importKey(key)
    rsakey = PKCS1_OAEP.new(rsakey)
    decrypted = rsakey.decrypt(b64decode(package))
    return decrypted

def chat_client():

    host = raw_input("Inserisci l'ip o l'hostname del server\n")
    port = input("Inserisci la porta su cui comunicare\n")
    nick=raw_input("Inserisci il tuo nickname\n")
    nick=nick + ": "
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(2)

    # connect to remote host
    try :
        s.connect((host, port))
    except :
        print 'Connessione non riuscita'
        sys.exit()

    print 'Connessione avvenuta con successo, adesso puoi iniziare a messaggiare'
    while 1:
        socket_list = [sys.stdin, s]

        # Get the list sockets which are readable
        read_sockets, write_sockets, error_sockets = select.select(socket_list , [], [])

        for sock in read_sockets:
            if sock == s:
                # incoming message from remote server, s
                data = sock.recv(4096)
                if not data :
                    print '\nDisconnesso dal server'
                    sys.exit()
                else :
                    #print data
                    b=decrypt_RSA(key_1, data)
                    sys.stdout.write(b)
            else :
                # user entered a message
                msg = sys.stdin.readline()
                msg = nick + msg
                a=encrypt_RSA(key_2, msg)
                s.send(a)
if __name__ == "__main__":

    sys.exit(chat_client())