import re
from ghidra.program.model.lang.OperandType import SCALAR, REGISTER
from ghidra.app.plugin.assembler import Assemblers
from ghidra.program.model.address import AddressSet
from common import *


ARGLOAD_SIGNATURE = b"\x4c\x8b\x40\x20\x48\x8b\x48\x18\x48\x8b\x50\x10\x48\x8b\x70\x08\x48\x8b\x38"

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

possible_indirect_callsites = findBytes(text, re.escape(ARGLOAD_SIGNATURE), 100, 1, True)

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

def f(a, b):
    n = 0
    while (a != 0):
        n += 1
        c = a
        a = (b - c) % 2**32
        b = c
    return b

base = f(0xbfb9c9a5, 0x2b495d78)

performed_at_least_one_fix = True
worklist = set(possible_indirect_callsites)
finished = set()

while performed_at_least_one_fix:
    performed_at_least_one_fix = False
    for callsite in worklist:
        post_argload_addr = callsite.add(len(ARGLOAD_SIGNATURE))
        disassemble(callsite)
        post_argload_ins = lst.getInstructionAt(post_argload_addr)
        print "Found possible callsite at %s" % (callsite)
        print post_argload_ins

        callsite_ins = lst.getInstructionAt(callsite)
        if not post_argload_ins:
            # already patched
            finished.add(callsite)
            continue

        if not (post_argload_ins.getMnemonicString() == "ADD" and post_argload_ins.getOperandType(1) == SCALAR):
            if post_argload_ins.getMnemonicString() == "MOV":
                print "Already patched by dynamic? %r: %r" % (post_argload_addr, post_argload_ins)
                finished.add(callsite)
            continue

        if callsite_ins.previous.getMnemonicString() != "MOV":
            continue

        add_ins = post_argload_ins
        call_ins = post_argload_ins.getNext()

        add_scalar = add_ins.getScalar(1).value

        call_reg = call_ins.getRegister(0)
        add_reg = call_reg.getChildRegisters()[0]

        mov_ins = callsite_ins.previous
        print mov_ins
        assert mov_ins.getRegister(0) == add_reg

        bytes = getBytes(callsite, post_argload_addr.subtract(callsite))
        clearListing(mov_ins.address, mov_ins.address.add(len(bytes)))
        setBytes(mov_ins.address, bytes)

        assemble_start_addr = mov_ins.address.add(len(bytes))

        call_target_addr = base + add_scalar
        disassemble(toAddr(call_target_addr))

        if not isRunningHeadless():
            setBackgroundColor(toAddr(call_target_addr), PURPLE)

        block = asm.assemble(assemble_start_addr, ["MOV %s, %d" % (add_reg, call_target_addr), "CALL %s" % (call_target_addr,)])

        setEOLComment(assemble_start_addr, "[static] single-target call patch")

        max_addr = block.getMaxAddress()
        nop_fill_sz = call_ins.next.address.subtract(max_addr) - 1
        assert nop_fill_sz >= 0
        asm.assemble(max_addr.add(1), ["NOP"] * nop_fill_sz)

        disassemble(mov_ins.address)
        finished.add(callsite)
        performed_at_least_one_fix = True
    worklist -= finished


possible_arr_callsites = findBytes(text, "\x42?\xff\x14....\x00", 500, 1, True)

to_remove = set()

for cs in possible_arr_callsites:
    if cs.add(1) in possible_arr_callsites:
        to_remove.add(cs.add(1))

for x in to_remove:
    possible_arr_callsites.remove(x)

for cs in possible_arr_callsites:
    ins = lst.getInstructionAt(cs)
    if str(ins).startswith("CALL qword ptr [0x410240"):
        asm.assemble(cs, ["CALL 0x4010b0", "NOP", "NOP", "NOP"])
    elif str(ins).startswith("CALL qword ptr [0x410280"):
        asm.assemble(cs, ["CALL 0x4010d0", "NOP", "NOP", "NOP"])
    else:
        continue
    setEOLComment(cs, "[static] str(n)cmp call")


if worklist:
    print "Warning: items left on worklist: %r" % (worklist)
