#!/usr/bin/python

# Simple client for the "count characters" protocol. Symmetric crypto
# with AES, to prevent sniffing. You can check with Wireshark that the
# data is unreadable.

import logging
import sys
import socket
from Crypto.Cipher import AES

PORT=4923
MAX=256
KEY="fubaaar12AZZ666."

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]
s = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
s.connect((server, PORT), )
outf = s.makefile('w')
inf = s.makefile('r')
coder = AES.new(KEY, AES.MODE_ECB)
# TODO: this code works for messages of length > 16?
if (len(message) % 16) != 0:
    n = 16 - (len(message) % 16)
    for i in range(0, n):
        message += '\0'
outf.write("%s" % coder.encrypt(message))
outf.close()
s.shutdown(socket.SHUT_WR)
response = inf.read(MAX)
inf.close()
logging.info("Response was %s" % response)
