import SocketServer
from hashlib import md5
import random
import math
import os
import struct

from itertools import *

kChunkSize = 128
kAddBit = 3

def print_banner(s):
  s.sendall("""Welcome to the Hash Game!\n""")

def print_prob(s, target_hash):
  s.sendall("send data with hash %s\n" % target_hash.encode('hex'))

def add(a, b):
  assert len(a) <= 8
  assert len(b) <= 8

  a = ("\x00" * (8-len(a))) + a
  b = ("\x00" * (8-len(b))) + b

  na = struct.unpack('>Q', a)[0]
  nb = struct.unpack('>Q', b)[0]

  res = struct.pack('>Q', na + nb)

  return res[-kAddBit:]

def xor(a, b):
  assert (len(a) == len(b))
  return ''.join(chr((ord(x) ^ ord(y)) % 256) for (x, y) in izip(a, b))

class MyTCPHandler(SocketServer.BaseRequestHandler):
  def handle(self):
    s = self.request

    print_banner(s)

    is_solved = False
    for stage in range(50):
      target_hash = os.urandom(16)
      print_prob(s, target_hash)

      size = int(s.recv(10).strip().rstrip())

      if not (0 <= size and size < 1024 * kChunkSize):
        s.sendall('Invalid data size\n')
        return

      data = s.recv(size)

      add_res = "\x00" * kAddBit
      xor_res = "\x00" * (16-kAddBit)
      for i in range(0, size, kChunkSize):
        digest = md5(data[i:i+kChunkSize]).digest()
        add_part = digest[:kAddBit]
        xor_part = digest[kAddBit:]

        add_res = add(add_res, add_part)
        xor_res = xor(xor_res, xor_part)

        hash_val = add_res + xor_res
        print 'hash = ' + hash_val.encode('hex')

      out_hash = add_res + xor_res

      if out_hash == target_hash:
        s.sendall("Success. Let's go to a next level\n")
      else:
        s.sendall("Fail... May be next time..\n")
        return;

    s.sendall("Congrats! you solve all problems.\n")
    s.sendall("FLAG IS '******' (without quote)\n")


if __name__ == "__main__":
  HOST, PORT = "0.0.0.0", 34343

  # Create the server, binding to localhost on port 9999
  SocketServer.TCPServer.allow_reuse_address = True
  server = SocketServer.TCPServer((HOST, PORT), MyTCPHandler)

  # Activate the server; this will keep running until you
  # interrupt the program with Ctrl-C
  server.serve_forever()