Il2CppInspectorRedux/Il2CppInspector.Common/Outputs/ScriptResources/Targets/IDA.py

# IDA-specific implementation
import ida_kernwin
import ida_name
import ida_idaapi
import ida_typeinf
import ida_bytes
import ida_nalt
import ida_ida
import ida_ua
import ida_segment
import ida_funcs
import ida_xref
import ida_pro

if ida_pro.IDA_SDK_VERSION > 770:
    import ida_srclang
    import ida_dirtree

    IDACLANG_AVAILABLE = True
    FOLDERS_AVAILABLE = True
    print("IDACLANG available")
else:
    IDACLANG_AVAILABLE = False
    FOLDERS_AVAILABLE = False

try:
    from typing import TYPE_CHECKING, Union

    if TYPE_CHECKING:
        from ..shared_base import (
            BaseStatusHandler,
            BaseDisassemblerInterface,
            ScriptContext,
        )
        import os
        from datetime import datetime
except ImportError:
    pass

TINFO_DEFINITE = (
    0x0001  # These only exist in idc for some reason, so we redefine it here
)
DEFAULT_TIL: "ida_typeinf.til_t" = None  # type: ignore


class IDADisassemblerInterface(BaseDisassemblerInterface):
    supports_fake_string_segment = True

    _status: BaseStatusHandler

    _type_cache: dict
    _folders: list[str]

    _function_dirtree: "ida_dirtree.dirtree_t"
    _cached_genflags: int
    _skip_function_creation: bool
    _is_32_bit: bool
    _fake_segments_base: int

    def __init__(self, status: BaseStatusHandler):
        self._status = status

        self._type_cache = {}
        self._folders = []

        self._cached_genflags = 0
        self._skip_function_creation = False
        self._is_32_bit = False
        self._fake_segments_base = 0

    def _get_type(self, type: str):
        if type not in self._type_cache:
            info = ida_typeinf.idc_parse_decl(DEFAULT_TIL, type, ida_typeinf.PT_RAWARGS)
            if info is None:
                print(f"Failed to create type {type}.")
                return None

            self._type_cache[type] = info[1:]

        return self._type_cache[type]

    def get_script_directory(self) -> str:
        return os.path.dirname(os.path.realpath(__file__))

    def on_start(self):
        # Disable autoanalysis
        self._cached_genflags = ida_ida.inf_get_genflags()
        ida_ida.inf_set_genflags(self._cached_genflags & ~ida_ida.INFFL_AUTO)

        # Unload type libraries we know to cause issues - like the c++ linux one
        PLATFORMS = ["x86", "x64", "arm", "arm64", "win7"]
        PROBLEMATIC_TYPELIBS = ["gnulnx", "mssdk64"]

        for lib in PROBLEMATIC_TYPELIBS:
            for platform in PLATFORMS:
                ida_typeinf.del_til(f"{lib}_{platform}")

        # Set name mangling to GCC 3.x and display demangled as default
        ida_ida.inf_set_demnames(ida_ida.DEMNAM_GCC3 | ida_ida.DEMNAM_NAME)

        self._status.update_step("Processing Types")

        if IDACLANG_AVAILABLE:
            header_path = os.path.join(
                self.get_script_directory(), "%TYPE_HEADER_RELATIVE_PATH%"
            )
            ida_srclang.set_parser_argv(
                "clang", "-target x86_64-pc-linux -x c++ -D_IDACLANG_=1"
            )  # -target required for 8.3+
            ida_srclang.parse_decls_with_parser("clang", None, header_path, True)
        else:
            original_macros = ida_typeinf.get_c_macros()
            ida_typeinf.set_c_macros(original_macros + ";_IDA_=1")
            ida_typeinf.idc_parse_types(
                os.path.join(
                    self.get_script_directory(), "%TYPE_HEADER_RELATIVE_PATH%"
                ),
                ida_typeinf.PT_FILE,
            )
            ida_typeinf.set_c_macros(original_macros)

        # Skip make_function on Windows GameAssembly.dll files due to them predefining all functions through pdata which makes the method very slow
        self._skip_function_creation = (
            ida_segment.get_segm_by_name(".pdata") is not None
        )
        if self._skip_function_creation:
            print(".pdata section found, skipping function boundaries")

        if FOLDERS_AVAILABLE:
            self._function_dirtree = ida_dirtree.get_std_dirtree(
                ida_dirtree.DIRTREE_FUNCS
            )

        self._is_32_bit = ida_ida.inf_is_32bit_exactly()

    def on_finish(self):
        ida_ida.inf_set_genflags(self._cached_genflags)

    def define_function(self, address: int, end: Union[int, None] = None):
        if self._skip_function_creation:
            return

        ida_bytes.del_items(
            address, ida_bytes.DELIT_SIMPLE, 12
        )  # Undefine x bytes which should hopefully be enough for the first instruction
        ida_ua.create_insn(address)  # Create instruction at start
        if not ida_funcs.add_func(
            address, end if end is not None else ida_idaapi.BADADDR
        ):  # This fails if the function doesn't start with an instruction
            print(
                f"failed to mark function {hex(address)}-{hex(end) if end is not None else '???'} as function"
            )

    def define_data_array(self, address: int, type: str, count: int):
        self.set_data_type(address, type)

        flags = ida_bytes.get_flags(address)
        if ida_bytes.is_struct(flags):
            opinfo = ida_nalt.opinfo_t()
            ida_bytes.get_opinfo(opinfo, address, 0, flags)
            entrySize = ida_bytes.get_data_elsize(address, flags, opinfo)
            tid = opinfo.tid
        else:
            entrySize = ida_bytes.get_item_size(address)
            tid = ida_idaapi.BADADDR

        ida_bytes.create_data(address, flags, count * entrySize, tid)

    def set_data_type(self, address: int, type: str):
        type += ";"

        info = self._get_type(type)
        if info is None:
            return

        if (
            ida_typeinf.apply_type(
                DEFAULT_TIL, info[0], info[1], address, TINFO_DEFINITE
            )
            is None
        ):
            print(f"set_type({hex(address)}, {type}); failed!")

    def set_function_type(self, address: int, type: str):
        self.set_data_type(address, type)

    def set_data_comment(self, address: int, cmt: str):
        ida_bytes.set_cmt(address, cmt, False)

    def set_function_comment(self, address: int, cmt: str):
        func = ida_funcs.get_func(address)
        if func is None:
            return

        ida_funcs.set_func_cmt(func, cmt, True)

    def set_data_name(self, address: int, name: str):
        ida_name.set_name(
            address, name, ida_name.SN_NOWARN | ida_name.SN_NOCHECK | ida_name.SN_FORCE
        )

    def set_function_name(self, address: int, name: str):
        self.set_data_name(address, name)

    def add_cross_reference(self, from_address: int, to_address: int):
        ida_xref.add_dref(from_address, to_address, ida_xref.XREF_USER | ida_xref.dr_I)

    def import_c_typedef(self, type_def: str):
        ida_typeinf.idc_parse_types(type_def, 0)

    # optional
    def add_function_to_group(self, address: int, group: str):
        if (
            not FOLDERS_AVAILABLE or ida_pro.IDA_SDK_VERSION < 930
        ):  # enable at your own risk on pre 9.3 - this is slow
            return

        if group not in self._folders:
            self._folders.append(group)
            self._function_dirtree.mkdir(group)

        name = ida_funcs.get_func_name(address)
        self._function_dirtree.rename(name, f"{group}/{name}")

    # only required if supports_fake_string_segment == True
    def create_fake_segment(self, name: str, size: int) -> int:
        start = ida_ida.inf_get_max_ea()
        end = start + size

        ida_segment.add_segm(0, start, end, name, "DATA")
        segment = ida_segment.get_segm_by_name(name)
        segment.bitness = 1 if self._is_32_bit else 2
        segment.perm = ida_segment.SEGPERM_READ
        segment.update()

        return start

    def write_string(self, address: int, value: str) -> int:
        encoded_string = value.encode() + b"\x00"
        string_length = len(encoded_string)
        ida_bytes.put_bytes(address, encoded_string)
        ida_bytes.create_strlit(address, string_length, ida_nalt.STRTYPE_C)
        return string_length

    def write_address(self, address: int, value: int):
        if self._is_32_bit:
            ida_bytes.put_dword(address, value)
        else:
            ida_bytes.put_qword(address, value)


# Status handler


class IDAStatusHandler(BaseStatusHandler):
    def __init__(self):
        self.step = "Initializing"
        self.max_items = 0
        self.current_items = 0
        self.start_time = datetime.now()
        self.step_start_time = self.start_time
        self.last_updated_time = datetime.min

    def initialize(self):
        ida_kernwin.show_wait_box("Processing")

    def update(self):
        if self.was_cancelled():
            raise RuntimeError("Cancelled script.")

        current_time = datetime.now()
        if 0.5 > (current_time - self.last_updated_time).total_seconds():
            return

        self.last_updated_time = current_time

        step_time = current_time - self.step_start_time
        total_time = current_time - self.start_time
        message = f"""
Running IL2CPP script.
Current Step: {self.step}
Progress: {self.current_items}/{self.max_items}
Elapsed: {step_time} ({total_time})
"""

        ida_kernwin.replace_wait_box(message)

    def update_step(self, step, max_items=0):
        print(step)

        self.step = step
        self.max_items = max_items
        self.current_items = 0
        self.step_start_time = datetime.now()
        self.last_updated_time = datetime.min
        self.update()

    def update_progress(self, new_progress=1):
        self.current_items += new_progress
        self.update()

    def was_cancelled(self):
        return ida_kernwin.user_cancelled()

    def shutdown(self):
        ida_kernwin.hide_wait_box()


status = IDAStatusHandler()
backend = IDADisassemblerInterface(status)
context = ScriptContext(backend, status)
context.process()