import phun
import capstone as cs
import unicorn as u
import hashlib
from unicorn.x86_const import *

#def get_relevant(
with open('./Choices') as f:
    data=f.read()
BASE = 0x400000
code = []

def print_bb(var,addr):
    if var == 0x831a13ee:
        addr = 0x40316A
    off = addr - BASE
#    print 'block:',hex(addr),'var:',hex(var)
    for ins in cs.disasm(data[off:off+1024],addr):
        if ins.mnemonic == 'jmp' or ins.mnemonic =='ret':
            break
#        print ins.mnemonic,ins.op_str
        code.append(ins.mnemonic + ' ' + ins.op_str)
    
cs= cs.Cs(cs.CS_ARCH_X86,cs.CS_MODE_64)
cs.detail = True

fd = open('./bb.x')
for line in fd:
    var = int(line.split(' ')[0],16)
    off = data.index(phun.p32(var)) + 6 + 4
    addr = BASE + off
    ins = cs.disasm(data[off:off+6],addr).next()
    print_bb(var,ins.operands[0].imm)
fd.close()

def hook_code(uc, addr, size, user_data):
    # ins=cs.disasm(str(uc.mem_read(addr,size)),addr).next()
    # print hex(int(addr)),ins.mnemonic,ins.op_str
    if addr == 0x100000e:
        print hex(mu.reg_read(UC_X86_REG_ECX))


DATA='45914878fa4c90081bf3a5a8d0464e3df54696bcb30e8820496955758b900738fbbf670bb0f0ebe10780ad5d1fc5788688e9be667991db82'.decode('hex')
ADDRESS = 0x1000000
DATA_A =  0x6000000
RBP = DATA_A + 0x200
bincode=phun.asm('\n'.join(code),'x64')
mu = u.Uc(u.UC_ARCH_X86, u.UC_MODE_64)
mu.mem_map(ADDRESS, 2 * 1024 * 1024)
mu.mem_map(DATA_A, 0x1000)
mu.reg_write(UC_X86_REG_RBP, RBP)
mu.mem_write(RBP-0x88,"\x00"*4)
mu.mem_write(RBP-0xd0,DATA)
mu.mem_write(RBP-0x10,phun.p64(DATA_A))
mu.mem_write(ADDRESS,bincode)
#mu.hook_add(u.UC_HOOK_CODE,hook_code)
try:
    mu.emu_start(ADDRESS, 1)
except:
    pass
flag=str(mu.mem_read(DATA_A,len(DATA)))
if hashlib.sha256(flag).hexdigest() == 'cc8fbb11d9a02ddbfb8f7468f08b7ef6d15a78d3c72b12cd3553d22700995ec8':
    print '[+] flag:',str(flag)