"""
Acha o RVA (offset relativo ao inicio do modulo) de uma funcao exportada, lendo o PE diretamente
do arquivo em disco (sem carregar/executar) -- necessario porque nosso python.exe e' 64-bit e o
kernel32.dll de 32-bit (SysWOW64) usado pelo AMo.exe e' um arquivo/mapeamento DIFERENTE do
kernel32.dll de 64-bit do nosso proprio processo (nao da pra so reusar GetProcAddress local).

Uso:
    python find_export.py <caminho_dll> <nome_funcao>
"""
import struct
import sys


def rva_to_offset(sections, rva):
    for va, vsize, praw, rawsize in sections:
        if va <= rva < va + max(vsize, rawsize):
            return praw + (rva - va)
    raise ValueError(f"RVA {rva:#x} nao mapeado em nenhuma secao")


def find_export_rva(dll_path, func_name):
    with open(dll_path, "rb") as f:
        data = f.read()

    e_lfanew = struct.unpack_from("<I", data, 0x3C)[0]
    pe_off = e_lfanew
    assert data[pe_off:pe_off + 4] == b"PE\x00\x00", "assinatura PE nao encontrada"

    file_header_off = pe_off + 4
    num_sections, = struct.unpack_from("<H", data, file_header_off + 2)
    size_opt_header, = struct.unpack_from("<H", data, file_header_off + 16)
    opt_header_off = file_header_off + 20

    magic, = struct.unpack_from("<H", data, opt_header_off)
    is_pe32_plus = (magic == 0x20B)
    data_dir_off = opt_header_off + (112 if is_pe32_plus else 96)

    export_rva, export_size = struct.unpack_from("<II", data, data_dir_off)

    section_table_off = opt_header_off + size_opt_header
    sections = []
    for i in range(num_sections):
        off = section_table_off + i * 40
        vsize, va, rawsize, praw = struct.unpack_from("<IIII", data, off + 8)
        sections.append((va, vsize, praw, rawsize))

    export_off = rva_to_offset(sections, export_rva)
    (characteristics, timestamp, majv, minv, name_rva, base,
     num_funcs, num_names, addr_funcs_rva, addr_names_rva, addr_ords_rva) = struct.unpack_from(
        "<IIHHIIIIIII", data, export_off)

    addr_funcs_off = rva_to_offset(sections, addr_funcs_rva)
    addr_names_off = rva_to_offset(sections, addr_names_rva)
    addr_ords_off = rva_to_offset(sections, addr_ords_rva)

    for i in range(num_names):
        name_rva_i, = struct.unpack_from("<I", data, addr_names_off + i * 4)
        name_off = rva_to_offset(sections, name_rva_i)
        end = data.index(b"\x00", name_off)
        name = data[name_off:end].decode("ascii")
        if name == func_name:
            ordinal, = struct.unpack_from("<H", data, addr_ords_off + i * 2)
            func_rva, = struct.unpack_from("<I", data, addr_funcs_off + ordinal * 4)
            return func_rva

    raise KeyError(f"{func_name} nao encontrado nos exports de {dll_path}")


if __name__ == "__main__":
    dll_path = sys.argv[1]
    func_name = sys.argv[2]
    rva = find_export_rva(dll_path, func_name)
    print(f"{func_name} RVA em {dll_path}: {rva:#x}")
