from ghidra.program.model.symbol import DataRefType
from ghidra.program.model.address import AddressSet
from ghidra.app.plugin.assembler import Assemblers


lst = currentProgram.getListing()
ref_mgr = currentProgram.getReferenceManager()
asm = Assemblers.getAssembler(currentProgram)

text = AddressSet(toAddr("4012d0"), toAddr("40eaf4"))


def get_retaddr_var(func):
    """ Get the function variable representing the return address. """
    vars = func.getAllVariables()
    for var in vars:
        if var.getMinAddress().getAddressSpace().name == u'stack' and var.getMinAddress().getOffset() == 0:
            # this is the return address
            return var
    return None


def handleRAModification(call, ra_read):
    """ NOP out instructions  """
    if not isRunningHeadless():
        setBackgroundColor(call, java.awt.Color(0x30fa50))
    call_ins = lst.getInstructionAt(call)
    ins = lst.getInstructionAt(ra_read.getFromAddress())
    print "[%r] Read: %r\n[%r] Add: %r\n[%r] Write: %r" % (ins.address, ins, ins.next.address, ins.next, ins.next.next.address, ins.next.next)
    add = ins.next
    assert add.mnemonicString == 'ADD'
    delta = add.getScalar(1).value
    asm.assemble(call_ins.getNext().address, ["NOP"] * delta)
    setEOLComment(call, "[static] return address patch")
    disassemble(call_ins.getNext().address.add(delta))


def isInPLT(addr):
    """ True if the address is within the .plt section, False otherwise. """
    return addr > toAddr("401020") and addr < toAddr("4012cf")


def run():
    close_relative_call = b"\xe8..[\xff\x00][\x00\xff]"
    calls = findBytes(text, close_relative_call, 500, 1, True)

    call_targets = {}


    for call in calls:
        disassemble(call)

        call_ins = lst.getInstructionAt(call)
        if not call_ins:
            continue

        if call_ins.next.mnemonicString == 'NOP':
            print "Already patched out by dynamic patcher? %r" % call
            continue

        target = call_ins.getAddress(0)

        if not text.contains(target):
            continue

        if not disassemble(target):
            continue

        createFunction(target, "direct_fn_%r" % target)

        call_targets[call] = getFunctionAt(target)


    analyzeChanges(currentProgram)


    fixed_up_something = True
    while fixed_up_something:
        fixed_up_something = False
        calls_to_remove = set()
        for call, func in call_targets.iteritems():
            ra_var = get_retaddr_var(func)
            if not ra_var:
                continue
            references = ref_mgr.getReferencesTo(ra_var)
            for ref in references:
                if ref.getReferenceType() == DataRefType.READ:
                    fixed_up_something = True
                    calls_to_remove.add(call)
                    handleRAModification(call, ref)
                    break

        for call in calls_to_remove:
            del call_targets[call]

run()
