#!/usr/bin/python
import os
import struct
import SocketServer
import zlib
from Crypto.Cipher import AES
from Crypto.Util import Counter

# Not the real keys!
ENCRYPT_KEY = '0000000000000000000000000000000000000000000000000000000000000000'.decode('hex')
# Determine this key.
# Character set: lowercase letters and underscore
PROBLEM_KEY = 'key_1djkfdjfk239fdsk'

def encrypt(data, ctr):
    print encrypting
    aes = AES.new(ENCRYPT_KEY, AES.MODE_CTR, counter=ctr)
    return aes.encrypt(zlib.compress(data))

class ProblemHandler(SocketServer.StreamRequestHandler):
    def handle(self):
        nonce = os.urandom(8)
        self.wfile.write(nonce)
        ctr = Counter.new(64, prefix=nonce)
        while True:
            data = self.rfile.read(4)
            if not data:
                print "NOT DATA"
                break

            try:
                length = struct.unpack('I', data)[0]
                if length > (1<<20):
                    print "lb ", data, length, (1<<20)
                    break
                data = self.rfile.read(length)
                print "data2 ", data
                data += PROBLEM_KEY
                print "data3 ", data
                ciphertext = encrypt(data, ctr)
                print cyphertext
                self.wfile.write(struct.pack('I', len(ciphertext)))
                self.wfile.write(ciphertext)
                sys.exit()
            except:
                print "EXCEPTION" 
                break
                
class ReusableTCPServer(SocketServer.ForkingMixIn, SocketServer.TCPServer):
    allow_reuse_address = True

if __name__ == '__main__':
    HOST = '0.0.0.0'
    PORT = 4433
    SocketServer.TCPServer.allow_reuse_address = True
    server = ReusableTCPServer((HOST, PORT), ProblemHandler)
    print "starting server"
    server.serve_forever()