#!/usr/bin/python

# Simple client for the "count characters" protocol. Authentication
# with RSA, to prevent in-transit modification.

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

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

logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s %(levelname)s %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S')
if len(sys.argv) != 3:
    logging.error("Usage: %s IPv6-address-of-server message" % sys.argv[0])
    sys.exit(-1)
server = sys.argv[1]
message = sys.argv[2]
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()
s = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
s.connect((server, PORT), )
outf = s.makefile('w')
inf = s.makefile('r')
signed_message = client_key.sign(message, None)[0]
outf.write("%s%s%li" % (message, SEPARATOR, signed_message))
outf.close()
s.shutdown(socket.SHUT_WR)
signed_response = inf.read(MAX)
(response, signature_str) = signed_response.split(SEPARATOR)
signature = long(signature_str)
result = server_key.verify(response, (signature, ))
inf.close()
logging.info("Response was %s and checked %s" % (response, result))