1import idc
2import idaapi
3import time
4
5from rpyc import classic
6c = classic.connect("127.0.0.1",port=18812)
7
8triton = c.modules.triton
9tast = c.modules['triton.ast']
10aexprs = c.modules['arybo.lib.mba_exprs']
11easm = c.modules['arybo.lib.exprs_asm']
12atools = c.modules['arybo.tools']
13triton.setArchitecture(triton.ARCH.X86_64)
14#triton.setAstRepresentationMode(triton.AST_REPRESENTATION.PYTHON)
15#triton.enableSymbolicOptimization(triton.OPTIMIZATION.ALIGNED_MEMORY, True)
16
17sym_rdi = triton.convertRegisterToSymbolicVariable(triton.REG.RDI, "rdi input")
18rdi = atools.tritonast2arybo(tast.variable(sym_rdi))
19print("[ ] %s = RDI" % str(sym_rdi))
20
21ea = idc.ScreenEA()
22func = idaapi.get_func(ea)
23
24pc = func.startEA
25print("[+] computing Triton AST for function starting at 0x%x, ending at 0x%x..." % (func.startEA, func.endEA))
26while pc < func.endEA-1:
27    inst = triton.Instruction()
28    opcode = idc.GetManyBytes(pc, idc.ItemSize(pc))
29    inst.setOpcodes(opcode)
30    inst.setAddress(pc)
31    triton.processing(inst)
32    pc = triton.getSymbolicRegisterValue(triton.REG.RIP)
33
34rax_ast = triton.buildSymbolicRegister(triton.REG.RAX)
35rax_ast = triton.getFullAst(rax_ast)
36rax_ast = triton.simplify(rax_ast, True)
37start = time.time()
38print("[+] computing Arybo representation...")
39e = atools.tritonast2arybo(rax_ast,use_exprs=True,use_esf=False)
40print("[+] got Arybo expression, evalute it...")
41e = aexprs.eval_expr(e,use_esf=False)
42end = time.time()
43diff = end-start
44print("[*] Arybo evaluation computed in %0.4fs" % diff)
45app = e.vectorial_decomp([rdi.v])
46exp = atools.identify(app,"rdi")
47print("[*] Identified expression: rax = %s" % aexprs.prettyprint(exp))
48asm = easm.asm_binary(exp, ("rax",64), {"rdi": ("rdi",64)}, "x86_64-unknown-unknwon")
49print("[*] Assembled expression: %s" % asm.encode("hex"))
50
51func_size = func.endEA-1-func.startEA
52if len(asm) > func_size:
53    printf("[-] Final assembly does not fit in the original function!")
54asm_nop = asm + "\x90"*(func_size-len(asm))
55func_start = int(func.startEA)
56func_end = int(func.endEA)
57idaapi.patch_many_bytes(func_start, asm_nop)
58idaapi.do_unknown_range(func_start, func_end, 0)
59idaapi.auto_make_code(func_start)
60idaapi.auto_make_proc(func_start)
61