from triton import *

def initCtx(t):
    t.setConcreteRegisterValue(t.registers.rsp,0x7fffffffde40)

def createSymVars(t):
    t.symbolizeRegister(t.registers.rbx,"SymRbx")
    rsp_val = t.getConcreteRegisterValue(t.registers.rsp)
    t.symbolizeMemory(MemoryAccess(rsp_val,CPUSIZE.QWORD),"SymRspAccess")
    t.symbolizeMemory(MemoryAccess(rsp_val,CPUSIZE.QWORD),"SymRspAdd2Access")

def init():
    t = TritonContext()
    t.setArchitecture(ARCH.X86_64)
    t.enableSymbolicEngine(True)
    t.enableTaintEngine(True)
    initCtx(t)
    createSymVars(t)
    return t

def load(t,entry,insns):
    pc = entry

    while True:
        insn = Instruction(pc,insns[pc])
        t.processing(insn)
        print(insn)
        if insn.isBranch():
            break
        pc = t.getRegisterAst(t.registers.rip).evaluate()
    return insn
def getConstraint(t,dstAddr):
    if type(dstAddr) != int:
        return None
    astctx = t.getAstContext()
    dstAddrBv = astctx.bv(dstAddr,64)
    rbx = t.getRegisterAst(t.registers.rbx)
    constraint = astctx.unroll(astctx.equal(rbx,dstAddrBv))
    return constraint

def getModelVars(m):
    MVars = dict()
    for i in range(len(m)):
        sm = m.get(i)
        if sm == None:
            continue
        sv = sm.getVariable()
        name = sv.getName()
        v = sm.getValue()
        MVars[name] = v
    return MVars

def main():

    insns = {0x400078:b"\x48\x0F\xC1\x1C\x24",0x40007d:b"\x48\x2B\x5C\x24\x02",0x400082:b"\x48\x0B\x1C\x24",0x400086:b"\xFF\xE3"}
    entry = 0x400078
    takenAddress = 0x40007f

    t = init()
    load(t,entry,insns)
    constraint = getConstraint(t,takenAddress)
    if not t.isSat(constraint):
        print("UNSAT")
        exit(-1)
    model = t.getModel(constraint)
    #modelVars = getModelVars(model)

    #ks = list(modelVars.keys())
    print("SAT\nmodel:",model,sep="\n\t")
    #for k in ks:
    #    print("%s=0x%x" % (k,modelVars[k]),end="\n\t")
    #print("\r")

if __name__ == "__main__":
    main()