import json
from common import *

asm = ghidra.app.plugin.assembler.Assemblers.getAssembler(currentProgram)

with open('./dyn_results.json', 'r') as f:
    data = json.load(f)

def find_ebx_move(start_ins):
    ins = start_ins

    while True:
        ins = ins.previous
        if str(ins) == "MOV EBX,EDX":
            return ins

lst = currentProgram.getListing()

for call, targets in data['call'].iteritems():
    call_addr = toAddr(call)
    call_ins = lst.getInstructionAt(call_addr)

    if not call_ins:
        disassemble(call_addr)

    call_ins = lst.getInstructionAt(call_addr)

    if len(targets) > 1:
        # Can't patch call target directly. We can, however, mark these calls
        # and their targets, and force disassembly

        print "Multi-target call at %s: %s" % (call, call_ins)
        disassemble(toAddr(call))
        setEOLComment(toAddr(call), "[dyn] multi-target call, calls %s" % (', '.join("{@address 0x%s}" % target for target in targets)))

        if not isRunningHeadless():
            setBackgroundColor(toAddr(call), (GREEN))

        for target in targets:
            disassemble(toAddr(target))
            if not isRunningHeadless():
                setBackgroundColor(toAddr(target), (PURPLE))
            setEOLComment(toAddr(target), "[dyn] target of multi-target call at {@address 0x%s}" % call)
    else:
        target = targets[0]
        if target.startswith('7f'):
            continue

        print "Single-target call at %s: %s" % (call_addr, str(call_ins))

        if str(call_ins) == 'CALL RBX':
            # Remove MOV EBX, EDX and replace ADD RBX, ...; CALL RBX with MOV RBX, TARGET; CALL TARGET; NOP
            ebx_move = find_ebx_move(call_ins)
            bytes = getBytes(ebx_move.getAddress(), call_ins.getAddress().subtract(ebx_move.getAddress()))
            bytes = bytes[2:-7]
            bytes.extend([-0x70, -0x70])
            clearListing(ebx_move.getAddress(), ebx_move.getAddress().add(len(bytes)))
            setBytes(ebx_move.getAddress(), bytes)
            disassemble(ebx_move.getAddress())

            add_addr = call_addr.subtract(9)
            asm.assemble(add_addr, ["MOV EBX, 0x%s" % target, "CALL 0x%s" % target, "NOP"])
            if not isRunningHeadless():
                setBackgroundColor(add_addr, (TURQUOISE))
            setEOLComment(add_addr, "[dyn] single-target call patch")
        elif str(call_ins).startswith("CALL qword ptr [0x410240"):
            # direct call to strncmp
            asm.assemble(call_addr, ["CALL 0x4010b0", "NOP", "NOP", "NOP"])
            setEOLComment(call_addr, "[dyn] indirect strncmp call")
        elif str(call_ins).startswith("CALL qword ptr [0x410280"):
            # direct call to strcmp
            asm.assemble(call_addr, ["CALL 0x4011d0", "NOP", "NOP", "NOP"])
            setEOLComment(call_addr, "[dyn] indirect strcmp call")
        else:
            if not isRunningHeadless():
                setBackgroundColor(call_addr, (ORANGE))
            setEOLComment(call_addr, "[dyn] unhandled single-target call")
