#!/usr/bin/env python
#-*- coding:utf-8 -*-

from sock import Sock
from hashlib import *

f = Sock("206.189.32.108 13579")
pref = f.read_line().strip().split()[-1].decode("hex")
print "got chal", pref.encode("hex")
for ch in xrange(2**30):
    h = sha256(pref + str(ch)).hexdigest()
    if h.startswith('00000'):
        break
else:
    print "fail"
    quit()

f.send_line(str(ch))

# modified https://github.com/eboda/mersenne-twister-recover
from MTRecover import MT19937Recover
from struct import unpack, pack

outputs = []
f.send("1\nasd\n" * 200)
for i in xrange(200):
    f.read_until("CHOOSE 1")
    f.read_until("give me a string: ")
    ct = f.read_line().strip().decode("hex")
    iv = ct[:16]
    outputs.extend(list(unpack(">4I", iv))[::-1])

mtr = MT19937Recover()
r2 = mtr.go(outputs)
for i in xrange(624):
    assert r2.prev() == outputs[623 - i]

# undo PoW
for i in xrange(2): print "%08x" % r2.prev()
print

# undo key rand
for i in xrange(4+4): r2.prev()

def pad(s):
    bs = 16
    padnum = bs - len(s) % bs
    return s + padnum * chr(padnum)

def to_string(num, max_len = 128):
    tmp = bin(num).lstrip('0b')[-max_len:].rjust(max_len, '0')
    return "".join(chr(int(tmp[i:i+8], 2)) for i in range(0, max_len, 8))

def gen_key(r):
    tmp1 = r.random()
    tmp2 = r.random()
    key = int(tmp1 * 2**128) | int(tmp2 * 2**75) | (0 & 0x3fffff)
    key = to_string(key)
    return key


f.send("1\nasd\n")
f.read_until("CHOOSE 1")
f.read_until("give me a string: ")
ct_asd = f.read_line().strip().decode("hex")

f.send("3\n")
f.read_until("CHOOSE 1")
f.read_until("> ")
ct_flag = f.read_line().strip().decode("hex")


r2 = r2.toPython()

k1 = gen_key(r2)
k2 = gen_key(r2)

from Crypto.Cipher import AES

iv = ct_asd[:16]
m = pad("asd")
c = ct_asd[16:32]


print "Stage 1"
tab = {}
for kval in xrange(2**22):
    if kval & 0x3ffff == 0: print hex(kval)
    t = unpack(">I", k1[-4:])[0]
    t |= kval
    key = k1[:-4] + pack(">I", t)

    a = AES.new(key, mode=AES.MODE_CBC, IV=iv)
    t = a.decrypt(c)

    tab[t] = key

print "Stage 2"
for kval in xrange(2**22):
    if kval & 0x3ffff == 0: print hex(kval)
    t = unpack(">I", k2[-4:])[0]
    t |= kval
    key = k2[:-4] + pack(">I", t)

    a = AES.new(key, mode=AES.MODE_CBC, IV=iv)
    t = a.encrypt(m)
    if t in tab:
        k1 = tab[t]
        k2 = key
        print "PWND", k1.encode("hex"), k2.encode("hex")
        break
else:
    print "FAIL"
    quit()

iv = ct_flag[:16]
c = ct_flag[16:]
a = AES.new(k1, mode=AES.MODE_CBC, IV=iv)
c = a.decrypt(c)
a = AES.new(k2, mode=AES.MODE_CBC, IV=iv)
c = a.decrypt(c)
print `c`
# PWND 1a873c3f0b75e891fb922765b0a98115 d33a05b242e90bbee4fe132377caa57c
# 'MeePwnCTF{DO_n0t_trust_anyth1ng}\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10\x10'