from triton import *
from capstone import *
from z3 import *
import os
import struct

entry_address = 0x400078
rsp_val = 0x7fffffffde40

def initRegs(t,SymCtx):
    global rsp_val
    t.concretizeAllRegister()
    SymCtx["symRbx"] = t.symbolizeRegister(t.registers.rbx)
    t.symbolizeRegister(t.registers.sf)
    t.symbolizeRegister(t.registers.of)
    t.setConcreteRegisterValue(t.registers.rsp,rsp_val)
    return SymCtx

def initMem(t,SymCtx):
    t.concretizeAllMemory()
    rspConcrete = t.getConcreteRegisterValue(t.registers.rsp)
    SymCtx["mem"] = MemoryAccess(rspConcrete,2)
    print("symbolize: ",SymCtx["mem"])
    SymCtx["symMemRsp2"] = t.symbolizeMemory(SymCtx["mem"])
    return SymCtx

def init(SymCtx):
    t = TritonContext()
    t.setArchitecture(ARCH.X86_64)
    t.enableSymbolicEngine(True)
    t.enableTaintEngine(True)
    SymCtx = initRegs(t,SymCtx)
    initMem(t,SymCtx)
    return t,SymCtx

def disasm(itext,address):
    cs = Cs(CS_ARCH_X86,CS_MODE_64)
    insns = dict()
    pc = address
    count = 0
    for insn in cs.disasm(itext,address):
        raw_insn = bytes(insn.bytes)
        insns[pc]=raw_insn
        size_insn = len(raw_insn)
        pc+=size_insn
        count+=1
    return insns,count

def get_decoded_insns(address):
    filename = "code.bin"
    f = open(filename,"rb")
    size = os.stat(filename).st_size
    itext = f.read(size)
    f.close()
    insns,count = disasm(itext,address)
    return insns,count

def load(t,SymCtx):
    global entry_address
    pc = entry_address
    insns,count = get_decoded_insns(address=pc)
    for i in range(count):
        insn = Instruction()
        insn.setAddress(pc)
        insn.setOpcode(insns[pc])
        supported = t.processing(insn)
        if not supported:
            return (False,insn)
        print(insn)
        if insn.isBranch():
            break
        AstNodeRip = t.getRegisterAst(t.registers.rip)
        pc = AstNodeRip.evaluate()
    return True,None,SymCtx

def getTrueModel(t,constraint):
    if t.isSat(constraint):
        return t.getModel(constraint)
    else:
        return None

def getFalseModel(t,constraint):
    astctx = t.getAstContext()
    constraint = astctx.lnot(constraint)
    if t.isSat(constraint):
        return t.getModel(constraint)
    else:
        return None

def extractVarsModel(t,model,SymCtx):
    symRbx = SymCtx["symRbx"]
    symMemRsp2 = SymCtx["symMemRsp2"]
    rbx_val = model[symRbx.getId()].getValue()
    memRsp2_val = model[symMemRsp2]
    return {"rbx":rbx_val,"mdword_rsp":MemRsp2_val}

def emulate_init(t,init_vals,SymCtx):
    global rsp_val
    t.reset()
    t.setArchitecture(ARCH.X86_64)
    t.enableSymbolicEngine(False)
    t.enableTaintEngine(False)
    rbx_val = init_vals["rbx"]
    mdword_rsp_le = init_vals['mdword_rsp']
    t.setConcreteRegisterValue(t.registers.rsp,rsp_val)
    t.setConcreteRegisterValue(t.registers.rbx,rbx_val)
    t.setConcreteMemoryValue(SymCtx["mem"],mdword_rsp_le)

def emulate(t,init_vals,SymCtx):
    global entry_address
    emulate_init(t,init_vals,SymCtx)
    insns,count = get_decoded_insns(entry_address)
    pc = entry_address
    for i in range(count):
        insn = Instruction()
        insn.setAddress(pc)
        insn.setOpcode(insns[pc])
        t.processing(insn)
        if insn.isBranch():
            is_branch_taken = insn.isConditionTaken()
            return is_branch_taken
        RipAstNode = t.getRegisterAst(t.registers.rip)
        pc = RipAstNode.evaluate()
    return None

def getConstraint(t):
    astctx = t.getAstContext()
    pcoRef = t.getPathPredicate()
    pcoFull = astctx.unroll(pcoRef)
    return pcoFull


def getBranchModels(t):
    astctx = t.getAstContext()
    constraint = getConstraint(t)
    model1 = getTrueModel(t,constraint)
    model2 = getFalseModel(t,constraint)
    return model1,model2

def solve(t,SymCtx):

    model1,model2 = getBranchModels(t,SymCtx)

    TVars = dict()
    FVars = dict()

    if model1 != None:
        TVars = extractVarsModel(t,model1,SymCtx)
        branch_taken = emulate(t,TVars,SymCtx)
        if branch_taken:
            keys = list(TVars.keys())
            print("\n[+] Complete Emulate for Branch Taken\n[*] Model representated in little endian:\n\t%s:%d,%s:%d" % (keys[0],TVars[keys[0]],keys[1],TVars[keys[1]]))
        else:
            print("[-] Not found model for Branch Taken")

    if model2 != None:
        FVars = extractVarsModel(model2,SymCtx)
        branch_taken = emulate(t,FVars,SymCtx)
        if branch_taken:
            keys = list(FVars.keys())
            print("\n[+] Complete Emulate for Branch No Taken\n[*] Model representated in little endian:\n\t%s:%d,%s:%d" % (keys[0],FVars[keys[0]],keys[1],FVars[keys[1]]))
    else:
        print("[-] Not found model for Branch No Taken")

def main():

    SymCtx = {
            "symRbx":None,
            "mem":None,
            "symMemRsp2":None
            }

    t,SymCtx = init(SymCtx)
    status,unsupInsn,SymCtx = load(t,SymCtx)
    if not status:
        print("error support ->",unsupInsn)
        os.exit(-1)
    solve(t,SymCtx)

if __name__ == "__main__":
    main()