import lief

import utils
from utils import Instructions, hex_address_to_memory_representation, is_32b, is_little_endian


def patch_direct_adress_call(pe: lief.PE.Binary, rva: int, instruction_offset: int):
    # We can manually patch the instruction here: FF 15 08 10 00 01 represents `call [0x01001080]`
    new_value = hex_address_to_memory_representation(
        hex(rva + pe.imagebase),
        is_32b(pe),
        is_little_endian(pe),
    )
    pe.patch_address(instruction_offset, Instructions.CALL_ADDR + new_value, lief.Binary.VA_TYPES.RVA)
    utils.print_debug(f"    Patched a call at addr {hex(pe.imagebase + instruction_offset)}")


def patch_direct_adress_jump(pe: lief.PE.Binary, rva: int, instruction_offset: int):
    # We can manually patch the instruction here: FF 15 08 10 00 01 represents `call [0x01001080]`
    new_value = hex_address_to_memory_representation(hex(rva + pe.imagebase), is_32b(pe), is_little_endian(pe))
    pe.patch_address(instruction_offset, Instructions.JUMP_ADDR + new_value, lief.Binary.VA_TYPES.RVA)
    utils.print_debug(f"    Patched a jump at addr {hex(pe.imagebase + instruction_offset)}")


def patch_instr_to_new_IAT_entry(pe: lief.PE.Binary, call: dict[str, str], rva: int):
    base = pe.imagebase
    instruction_offset = int(call["adress"], 16) - base
    memview = pe.get_content_from_virtual_address(instruction_offset, 2)
    if [memview[0], memview[1]] == Instructions.CALL_ADDR:
        patch_direct_adress_call(pe, rva, instruction_offset)
    elif [memview[0], memview[1]] == Instructions.JUMP_ADDR:
        patch_direct_adress_jump(pe, rva, instruction_offset)


def patch_addr_found_in_mem(pe: lief.PE.Binary, rva: int, old_addr: str):
    is_32 = pe.abstract.header.is_32
    little_endian = pe.abstract.header.endianness == lief.Header.ENDIANNESS.LITTLE
    # scan memory for reference to old addr
    old_addr_mem_repr = hex_address_to_memory_representation(old_addr, is_32b(pe), is_little_endian(pe))
    new_addr = hex_address_to_memory_representation(hex(rva + pe.imagebase), is_32, little_endian)
    found_ref_addr = []
    found_xref_addr = []
    for section in pe.sections:
        for i in range(len(section.content)):
            found = True
            for j in range(len(old_addr_mem_repr)):
                if i + j >= len(section.content) or section.content[i + j] != old_addr_mem_repr[j]:
                    found = False
                    break
            if found:
                ref_addr = hex_address_to_memory_representation(
                    hex(
                        section.virtual_address + i + pe.imagebase,
                    ),
                    is_32,
                    little_endian,
                )
                found_ref_addr.append(ref_addr)

    for section in pe.sections:
        for ref_addr in found_ref_addr:
            for k in range(len(section.content) - len(ref_addr)):
                foundxref = True
                for L in range(len(ref_addr)):
                    if section.content[k + L] != ref_addr[L]:
                        foundxref = False
                        break
                if foundxref:
                    found_xref_addr.append(section.virtual_address + k)
    for addr in found_xref_addr:
        pe.patch_address(addr, new_addr, lief.Binary.VA_TYPES.RVA)
        utils.print_debug(f"    Patched an xref to old IAT at {hex(pe.imagebase + addr)}")