MEMCHK64_DYNAMIC_EXEC.py

# IDAPython 7.x / 8.x
# Finds calls to RE_resolve_kernel_api, extracts obfuscated string from RCX,
# Atbash-deobfuscates it, comments the LEA/CALL, then renames the stored
# resolved API pointer, e.g. unk_140005098 -> ptr_RtlInitUnicodeString.

import re

import idaapi
import idautils
import idc
import ida_bytes
import ida_funcs
import ida_name
import ida_nalt
import ida_ua


TARGET_NAME = "RE_resolve_kernel_api"


def atbash(s: str) -> str:
    out = []
    for c in s:
        if "a" <= c <= "z":
            out.append(chr(ord("z") - (ord(c) - ord("a"))))
        elif "A" <= c <= "Z":
            out.append(chr(ord("Z") - (ord(c) - ord("A"))))
        else:
            out.append(c)
    return "".join(out)


def sanitize_name(s: str) -> str:
    s = re.sub(r"[^0-9A-Za-z_]", "_", s).strip("_")
    if not s:
        s = "resolved_api"
    if s[0].isdigit():
        s = "_" + s
    return s


def get_c_string(ea: int):
    raw = ida_bytes.get_strlit_contents(ea, -1, ida_nalt.STRTYPE_C)
    if raw is None:
        raw = idc.get_strlit_contents(ea, -1, idc.STRTYPE_C)

    if raw is None:
        return None

    if isinstance(raw, bytes):
        return raw.decode("latin-1", errors="replace")

    return str(raw)


def reg_name(op, size=8):
    if op.type != ida_ua.o_reg:
        return None
    return idaapi.get_reg_name(op.reg, size).lower()


def is_reg(op, wanted):
    name = reg_name(op, 8)
    return name == wanted.lower()


def resolve_lea_string_ea(insn):
    """
    Handles:
        lea rcx, aString
        lea rcx, [rip+disp]
    """
    if insn.itype != idaapi.NN_lea:
        return None

    if not is_reg(insn.ops[0], "rcx"):
        return None

    op = insn.ops[1]

    if op.type in (ida_ua.o_mem, ida_ua.o_imm, ida_ua.o_far, ida_ua.o_near, ida_ua.o_displ):
        return op.addr

    return None


def find_rcx_string_before_call(call_ea: int, max_back: int = 25):
    """
    Walks backward in the same function looking for:
        lea rcx, obfuscated_string

    Stops if RCX is overwritten before finding the LEA.
    """
    func = ida_funcs.get_func(call_ea)
    min_ea = func.start_ea if func else 0

    ea = call_ea
    for _ in range(max_back):
        ea = idc.prev_head(ea, min_ea)
        if ea == idc.BADADDR or ea < min_ea:
            break

        insn = ida_ua.insn_t()
        if ida_ua.decode_insn(insn, ea) <= 0:
            continue

        str_ea = resolve_lea_string_ea(insn)
        if str_ea:
            s = get_c_string(str_ea)
            if s:
                return ea, str_ea, s

        # Avoid using an older unrelated RCX argument.
        if insn.ops[0].type == ida_ua.o_reg:
            dst = idaapi.get_reg_name(insn.ops[0].reg, 8).lower()
            if dst in ("rcx", "ecx", "cx", "cl"):
                break

    return None, None, None


def find_store_after_call(call_ea: int, max_forward: int = 12):
    """
    Looks forward for:
        mov qword ptr cs:unk_xxxxxxxx, rax

    Returns:
        store_instruction_ea, destination_memory_ea
    """
    func = ida_funcs.get_func(call_ea)
    max_ea = func.end_ea if func else idc.BADADDR

    ea = call_ea
    for _ in range(max_forward):
        ea = idc.next_head(ea, max_ea)
        if ea == idc.BADADDR or ea >= max_ea:
            break

        insn = ida_ua.insn_t()
        if ida_ua.decode_insn(insn, ea) <= 0:
            continue

        if idc.print_insn_mnem(ea).lower() != "mov":
            continue

        dst = insn.ops[0]
        src = insn.ops[1]

        src_is_rax = src.type == ida_ua.o_reg and idaapi.get_reg_name(src.reg, 8).lower() == "rax"
        dst_is_mem = dst.type in (ida_ua.o_mem, ida_ua.o_displ, ida_ua.o_far, ida_ua.o_near)

        if src_is_rax and dst_is_mem:
            return ea, dst.addr

        # If RAX gets clobbered before a store, stop.
        if dst.type == ida_ua.o_reg:
            dst_name = idaapi.get_reg_name(dst.reg, 8).lower()
            if dst_name in ("rax", "eax", "ax", "al"):
                break

    return None, None


def add_comment(ea: int, text: str):
    old = idc.get_cmt(ea, 0) or ""
    if text in old:
        return

    new = text if not old else old + "\n" + text
    idc.set_cmt(ea, new, 0)


def set_unique_name(ea: int, base_name: str):
    """
    Renames EA to base_name. If taken, appends _1, _2, etc.
    """
    name = base_name

    for i in range(1000):
        candidate = name if i == 0 else f"{name}_{i}"

        existing = idc.get_name_ea_simple(candidate)
        if existing not in (idc.BADADDR, ea):
            continue

        if ida_name.set_name(ea, candidate, ida_name.SN_CHECK):
            return candidate

    return None


def main():
    target_ea = idc.get_name_ea_simple(TARGET_NAME)
    if target_ea == idc.BADADDR:
        print(f"[-] Could not find {TARGET_NAME}")
        return

    annotated = 0
    renamed = 0

    for xref in idautils.XrefsTo(target_ea):
        call_ea = xref.frm

        if idc.print_insn_mnem(call_ea).lower() != "call":
            continue

        lea_ea, str_ea, obf = find_rcx_string_before_call(call_ea)
        if not obf:
            print(f"[!] No RCX string found before call at {call_ea:#x}")
            continue

        deobf = atbash(obf)
        comment = f'deobfuscated: "{deobf}"'

        add_comment(lea_ea, comment)
        add_comment(call_ea, comment)

        store_ea, ptr_ea = find_store_after_call(call_ea)
        if ptr_ea:
            base_name = "ptr_" + sanitize_name(deobf)
            final_name = set_unique_name(ptr_ea, base_name)

            if final_name:
                add_comment(store_ea, f"stores resolved API pointer: {final_name}")
                print(f"[+] {ptr_ea:#x} renamed to {final_name}")
                renamed += 1
            else:
                print(f"[!] Failed to rename {ptr_ea:#x} to {base_name}")
        else:
            print(f"[!] No post-call RAX store found after {call_ea:#x}")

        print(f"[+] {call_ea:#x}: {obf!r} -> {deobf!r}")
        annotated += 1

    print(f"[+] Done. Annotated {annotated} resolver call(s), renamed {renamed} pointer(s).")



main()