#!/usr/bin/env python

# Simple server for the "count characters" protocol. Asymmetric crypto
# with RSA, to prevent sniffing.

import SocketServer
import socket
import logging
from Crypto.PublicKey import RSA
import cPickle as pickle

PORT=4923
client_key_filename = "./client-public.key"
server_key_filename = "./server-full.key"

class RequestHandler(SocketServer.StreamRequestHandler):

    def handle(self):
        logging.info("Client connection")
        text = "DUMMY"
        encrypted_data = ""
        size = 0
        while text != "":
            encrypted_text = self.rfile.read(1)
            if encrypted_text == "":
                break
            encrypted_data += encrypted_text
        data = self.server.server_key.decrypt(encrypted_data)
        # TODO: handle UnicodeDecodeError
        text = data.decode("UTF-8")
        size = len(text)
        response = (u"You sent %i characters" % size).encode("UTF-8")
        self.wfile.write(self.server.client_key.\
                         encrypt(response, None)[0])

class Server(SocketServer.ThreadingMixIn, SocketServer.TCPServer): 

    def __init__(self, address, handler, mykey, client_key):
        # Do not use CBC, since each client will start from scratch,
        # but with an existing server. Security-wise, it is not a good
        # idea, we should use CBC and, to do so, the decoder should be
        # per-connection, not global.
        self.server_key = mykey
        self.client_key = client_key
        SocketServer.TCPServer.__init__(self, address, handler)
        logging.info("Starting server")

logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s %(levelname)s %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S')
server_key_file = open(server_key_filename, 'r')
server_key = pickle.load(server_key_file)
server_key_file.close()
client_key_file = open(client_key_filename, 'r')
client_key = pickle.load(client_key_file)
client_key_file.close()
Server.address_family = socket.AF_INET6
Server.allow_reuse_address = True
myserver = Server(("", PORT), RequestHandler, server_key, client_key)
try:
    myserver.serve_forever()
except KeyboardInterrupt:
    logging.info("Server stopped")