diff --git a/pyproject.toml b/pyproject.toml index 5158749..b0be806 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,4 +66,5 @@ testpaths = ["tests"] addopts = "--import-mode=importlib -s" markers = [ "idalib: tests that require an activated idalib (IDA Pro headless). Run with `pytest -m idalib`. In CI these run on the `ida` self-hosted runner image (IDA + idalib preinstalled, licensing provided by the runner environment).", + "ida_ui: tests that launch the real IDA GUI to observe widget behavior. Opt-in: set REAI_UI_TESTS=1 and ensure no interactive IDA instance is running (needs the license seat and a display).", ] diff --git a/reai_toolkit/app/components/dialogs/matching_dialog.py b/reai_toolkit/app/components/dialogs/matching_dialog.py index 2243450..9c40754 100644 --- a/reai_toolkit/app/components/dialogs/matching_dialog.py +++ b/reai_toolkit/app/components/dialogs/matching_dialog.py @@ -1240,7 +1240,7 @@ def import_data_types(self): remote_function_id: int = self.local_func_id_to_remote_func_id[local_function_id] selected_matches[remote_function_id] = self.matched_func_to_original_ea[remote_function_id] - self.data_types_service.import_data_types(selected_matches) + self.data_types_service.import_data_types_async(selected_matches) # ===================================================================== # (Optional) page-switch helpers diff --git a/reai_toolkit/app/services/data_types/data_types_service.py b/reai_toolkit/app/services/data_types/data_types_service.py index 022924c..d22b831 100644 --- a/reai_toolkit/app/services/data_types/data_types_service.py +++ b/reai_toolkit/app/services/data_types/data_types_service.py @@ -1,3 +1,4 @@ +import threading from dataclasses import dataclass, field from revengai import ApiException, Configuration @@ -30,6 +31,14 @@ class ImportDataTypesService(IThreadService): def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration) -> None: super().__init__(netstore_service=netstore_service, sdk_config=sdk_config) + def import_data_types_async(self, matches: dict[int, int]) -> None: + self.start_worker(target=self._import_worker, args=(matches,)) + + def _import_worker(self, _: threading.Event, matches: dict[int, int]) -> None: + result: DataTypesImportResult = self.import_data_types(matches) + if result.error: + logger.error(f"RevEng.AI: {result.error}") + def import_data_types(self, matches: dict[int, int]) -> DataTypesImportResult: if len(matches) == 0: return DataTypesImportResult() diff --git a/reai_toolkit/app/services/variable_sync/variable_sync_service.py b/reai_toolkit/app/services/variable_sync/variable_sync_service.py index c439c16..93aa089 100644 --- a/reai_toolkit/app/services/variable_sync/variable_sync_service.py +++ b/reai_toolkit/app/services/variable_sync/variable_sync_service.py @@ -4,7 +4,8 @@ from typing import Optional, Tuple from libbs.artifacts import Enum, FunctionHeader, StackVariable, Struct, Typedef -from libbs.decompilers.ida.compat import execute_read +from libbs.decompilers.ida.compat import DummyIDACodeView, execute_read +from libbs.decompilers.ida.compat import function as read_ida_function from loguru import logger from revengai import ( @@ -29,7 +30,20 @@ @execute_read def _read_decompiler_function(deci, func_addr: int): - return deci.functions.get(func_addr) + lowered: int = deci.art_lifter.lower_addr(func_addr) + try: + if deci.decompiler_available: + code_view = DummyIDACodeView(lowered) + if code_view.cfunc is None: + return None + func = read_ida_function(lowered, ida_code_view=code_view) + else: + func = read_ida_function(lowered, decompiler_available=False) + except Exception as e: + logger.debug(f"RevEng.AI: could not read function at 0x{lowered:x}: {e}") + return None + + return deci.art_lifter.lift(func) if func is not None else None @execute_read diff --git a/reai_toolkit/app/transformations/import_data_types.py b/reai_toolkit/app/transformations/import_data_types.py index 0918428..59cd2b4 100644 --- a/reai_toolkit/app/transformations/import_data_types.py +++ b/reai_toolkit/app/transformations/import_data_types.py @@ -1,21 +1,25 @@ from typing import cast +import ida_funcs +import ida_typeinf +import idaapi + import libbs.artifacts from libbs.api import DecompilerInterface -from libbs.decompilers.ida.compat import execute_ui, convert_type_str_to_ida_type +from libbs.decompilers.ida.compat import execute_write, convert_type_str_to_ida_type from loguru import logger from revengai import ( - FunctionArgument, Enumeration, + FunctionArgument, FunctionDataTypesList, - FunctionHeader, - V2FunctionInfoFuncDepsInner, - FunctionInfo, + FunctionDataTypesListItem, FunctionType, Structure, TypeDefinition, ) +APPLY_CHUNK_SIZE = 50 + class TaggedDependency: def __init__(self, dependency: Structure | Enumeration | TypeDefinition) -> None: @@ -31,69 +35,144 @@ class ImportDataTypes: def __init__(self) -> None: self.deci: DecompilerInterface - @execute_ui def execute(self, functions: FunctionDataTypesList, matched_function_mapping: dict[int, int] = {}) -> set[int]: - self.deci = DecompilerInterface.discover(force_decompiler="ida") # type: ignore - lookup: dict[str, TaggedDependency] = {} + items: list[FunctionDataTypesListItem] = [ + item for item in functions.items if item.data_types is not None + ] + if not items: + return set() + + # Track processed dependencies to prevent duplicate imports. + # Without this: + # - Shared dependencies get re-processed, breaking references (shows as invalid ordinals in IDA) + # - Cannot resolve subdependencies (e.g. struct fields that reference other imported types) + lookup: dict[str, TaggedDependency] = self._build_lookup(items) + if lookup: + self._apply_dependencies(lookup) + failed: set[int] = set() + total: int = len(items) + for start in range(0, total, APPLY_CHUNK_SIZE): + chunk: list[FunctionDataTypesListItem] = items[start:start + APPLY_CHUNK_SIZE] + failed |= self._apply_chunk(chunk, matched_function_mapping) + logger.info( + f"RevEng.AI: applied data types to {min(start + APPLY_CHUNK_SIZE, total)}/{total} functions" + ) - for function in functions.items: - data_types: FunctionInfo | None = function.data_types + return failed - if data_types is None: - continue - - # Track processed dependencies to prevent duplicate imports. - # Without this: - # - Shared dependencies get re-processed, breaking references (shows as invalid ordinals in IDA) - # - Cannot resolve subdependencies (e.g. struct fields that reference other imported types) - for dep in data_types.func_deps: + def _build_lookup(self, items: list[FunctionDataTypesListItem]) -> dict[str, TaggedDependency]: + lookup: dict[str, TaggedDependency] = {} + for item in items: + for dep in item.data_types.func_deps or []: if dep.actual_instance is None: continue if dep.actual_instance.name not in lookup: - lookup.update({dep.actual_instance.name: TaggedDependency(dep.actual_instance)}) # type: ignore + lookup[dep.actual_instance.name] = TaggedDependency(dep.actual_instance) # type: ignore - dependency: V2FunctionInfoFuncDepsInner - for dependency in data_types.func_deps: - if dependency.actual_instance is None: - continue + return lookup + + @execute_write + def _apply_dependencies(self, lookup: dict[str, TaggedDependency]) -> None: + self.deci = DecompilerInterface.discover(force_decompiler="ida") # type: ignore + for tagged_dependency in lookup.values(): + try: + self.process_dependency(tagged_dependency, lookup) + except Exception as e: + logger.warning( + f"RevEng.AI: skipped dependency {tagged_dependency.name!r}: {e!r}" + ) + tagged_dependency.processed = True + + @execute_write + def _apply_chunk( + self, chunk: list[FunctionDataTypesListItem], matched_function_mapping: dict[int, int] + ) -> set[int]: + failed: set[int] = set() + for item in chunk: + func: FunctionType | None = item.data_types.func_types + if func is None: + continue - tagged_dependency: TaggedDependency | None = lookup.get(dependency.actual_instance.name) - if tagged_dependency: - try: - self.process_dependency(tagged_dependency, lookup) - except Exception as e: - logger.warning( - f"RevEng.AI: skipped dependency {tagged_dependency.name!r}: {e!r}" - ) - tagged_dependency.processed = True - - func: FunctionType | None = data_types.func_types - if func: - try: - if matched_function_mapping: - ea: int = matched_function_mapping[function.function_id] - else: - ea: int = func.addr - - if not self.update_function(func, ea) or self._function_types_unparseable(func): - failed.add(function.function_id) - except Exception as e: - logger.warning( - f"RevEng.AI: skipped data types for function {function.function_id}: {e!r}" - ) - failed.add(function.function_id) + try: + if matched_function_mapping: + ea: int = matched_function_mapping[item.function_id] + else: + ea: int = func.addr + + if not self.apply_function_type(func, ea): + failed.add(item.function_id) + except Exception as e: + logger.warning( + f"RevEng.AI: skipped data types for function {item.function_id}: {e!r}" + ) + failed.add(item.function_id) return failed + def apply_function_type(self, func: FunctionType, ea: int) -> bool: + if ida_funcs.get_func(ea) is None: + logger.warning(f"failed to update function: {func.name} at 0x{ea:x}") + return False + + args: list[FunctionArgument] = sorted(func.header.args.values(), key=lambda a: a.offset) + arg_types: list[tuple[str, ida_typeinf.tinfo_t]] = [] + for arg in args: + arg_tif = convert_type_str_to_ida_type(self.normalise_type(arg.type)) if arg.type else None + if arg_tif is None: + return False + arg_types.append((arg.name, arg_tif)) + + details: ida_typeinf.func_type_data_t = self._current_func_details(ea) + + ret_type: str = self.normalise_type(func.header.type) if func.header.type else "" + if ret_type: + ret_tif = convert_type_str_to_ida_type(ret_type) + if ret_tif is None: + return False + details.rettype = ret_tif + elif details.rettype.empty(): + details.rettype = convert_type_str_to_ida_type("void") + + if arg_types: + details.clear() + for name, arg_tif in arg_types: + funcarg = ida_typeinf.funcarg_t() + funcarg.name = name + funcarg.type = arg_tif + details.push_back(funcarg) + + proto = ida_typeinf.tinfo_t() + if not proto.create_func(details): + return False + + return bool(ida_typeinf.apply_tinfo(ea, proto, ida_typeinf.TINFO_DEFINITE)) + + @staticmethod + def _current_func_details(ea: int) -> "ida_typeinf.func_type_data_t": + details = ida_typeinf.func_type_data_t() + existing = ida_typeinf.tinfo_t() + if idaapi.get_tinfo(existing, ea) and existing.is_func() and existing.get_func_details(details): + return details + + if ( + ida_typeinf.guess_tinfo(existing, ea) != ida_typeinf.GUESS_FUNC_FAILED + and existing.is_func() + and existing.get_func_details(details) + ): + return details + + details = ida_typeinf.func_type_data_t() + details.cc = ida_typeinf.CM_CC_UNKNOWN + return details def process_dependency( self, tagged_dependency: TaggedDependency, lookup: dict[str, TaggedDependency] ) -> None: if tagged_dependency.processed: return - + dependency: Structure | Enumeration | TypeDefinition = tagged_dependency.dependency match dependency: case Structure(): @@ -110,7 +189,7 @@ def process_dependency( def update_struct(self, imported_struct: Structure, lookup: dict[str, TaggedDependency]) -> None: if imported_struct.size is None: return - + for member in imported_struct.members.values(): subdependency: TaggedDependency | None = lookup.get(member.type) if subdependency: @@ -134,56 +213,6 @@ def update_typedef(self, imported_typedef: TypeDefinition, lookup: dict[str, Tag name=imported_typedef.name, type_=normalized_type ) - def update_function(self, func: FunctionType, ea: int) -> bool: - base_address: int = self.deci.binary_base_addr - rva: int = ea - base_address - - target_func: libbs.artifacts.Function | None = self.deci.functions.get(rva) # type: ignore - if target_func is None: - logger.warning(f"failed to update function: {func.name} at rva: 0x{rva:0x}") - return False - - target_func.name = func.name - target_func.size = func.size - target_func.type = func.type - - # Check the target function has a header. - if target_func.header: - self.update_header(func.header, target_func) - - self.deci.functions[rva] = target_func - return True - - def _function_types_unparseable(self, func: FunctionType) -> bool: - if not (func.header and func.header.args): - return False - for arg in func.header.args.values(): - type_str: str = self.normalise_type(arg.type) if arg.type else "" - if type_str and convert_type_str_to_ida_type(type_str) is None: - return True - return False - - def update_header( - self, imported_header: FunctionHeader, target_function: libbs.artifacts.Function - ) -> None: - if target_function.header is None: - return - - target_function.header.name = imported_header.name - target_function.header.type = self.normalise_type(imported_header.type) - self.update_function_arguments(imported_header.args, target_function) - - def update_function_arguments( - self, imported_args: dict[str, FunctionArgument], target_function: libbs.artifacts.Function - ) -> None: - if target_function.header is None: - return - - for arg in imported_args.values(): - arg.type = self.normalise_type(arg.type) - - target_function.header.args = {v.offset: v for v in imported_args.values()} - @staticmethod def normalise_type(data_type: str) -> str: # When we obtain a type from DWARF information, it often looks something like `DWARF/stdint-uintn.h::uint32_t` diff --git a/tests/idalib/data_types/test_import.py b/tests/idalib/data_types/test_import.py new file mode 100644 index 0000000..f8c7ca4 --- /dev/null +++ b/tests/idalib/data_types/test_import.py @@ -0,0 +1,123 @@ +import pytest + +import ida_typeinf +import idaapi +import idc +from revengai import ( + FunctionArgument, + FunctionDataTypesList, + FunctionDataTypesListItem, + FunctionHeader, + FunctionInfo, + FunctionType, +) + +from reai_toolkit.app.transformations.import_data_types import ImportDataTypes + +pytestmark = pytest.mark.idalib + + +@pytest.fixture +def func_ea(loaded_binary): + ea = idc.get_name_ea_simple("sub_401020") + assert ea != idaapi.BADADDR + original = ida_typeinf.tinfo_t() + had_type = idaapi.get_tinfo(original, ea) + yield ea + if had_type: + ida_typeinf.apply_tinfo(ea, original, ida_typeinf.TINFO_DEFINITE) + else: + idc.SetType(ea, "") + + +def _arg(offset: int, name: str, type_str: str) -> FunctionArgument: + return FunctionArgument(name=name, offset=offset, size=8, type=type_str) + + +def _func_types(ea: int, ret: str = "int", args: tuple = ()) -> FunctionType: + return FunctionType( + addr=ea, + header=FunctionHeader( + addr=ea, + args={hex(arg.offset): arg for arg in args}, + name="reai_test_func", + type=ret, + ), + name="reai_test_func", + size=16, + type=ret, + ) + + +def _item(function_id: int, func_types: FunctionType) -> FunctionDataTypesListItem: + return FunctionDataTypesListItem.model_construct( + function_id=function_id, + data_types=FunctionInfo.model_construct(func_deps=[], func_types=func_types), + ) + + +def test_apply_function_type_sets_prototype_with_named_args(func_ea): + idt = ImportDataTypes() + func = _func_types(func_ea, ret="int", args=(_arg(0, "count", "int"), _arg(1, "buf", "char *"))) + + assert idt.apply_function_type(func, func_ea) is True + + printed = ida_typeinf.print_type(func_ea, ida_typeinf.PRTYPE_1LINE) + assert "count" in printed + assert "buf" in printed + assert "char *" in printed + + +def test_apply_function_type_normalises_dwarf_types(func_ea): + idt = ImportDataTypes() + func = _func_types(func_ea, ret="uchar", args=(_arg(0, "n", "DWARF/stdint.h::qword"),)) + + assert idt.apply_function_type(func, func_ea) is True + + printed = ida_typeinf.print_type(func_ea, ida_typeinf.PRTYPE_1LINE) + assert "unsigned __int8" in printed + assert "unsigned __int64 n" in printed + + +def test_apply_function_type_keeps_args_when_remote_has_none(func_ea): + idt = ImportDataTypes() + assert idt.apply_function_type( + _func_types(func_ea, ret="int", args=(_arg(0, "count", "int"),)), func_ea + ) + + assert idt.apply_function_type(_func_types(func_ea, ret="void", args=()), func_ea) is True + + printed = ida_typeinf.print_type(func_ea, ida_typeinf.PRTYPE_1LINE) + assert "void" in printed + assert "count" in printed + + +def test_apply_function_type_rejects_unparseable_arg(func_ea): + idt = ImportDataTypes() + func = _func_types(func_ea, args=(_arg(0, "x", "totally bogus type!!"),)) + + assert idt.apply_function_type(func, func_ea) is False + + +def test_apply_function_type_missing_function(loaded_binary): + idt = ImportDataTypes() + + assert idt.apply_function_type(_func_types(0x1), 0x1) is False + + +def test_execute_applies_via_mapping_and_reports_failures(func_ea): + items = [ + _item(1, _func_types(func_ea, ret="int", args=(_arg(0, "count", "int"),))), + _item(2, _func_types(func_ea, args=(_arg(0, "x", "totally bogus type!!"),))), + _item(3, _func_types(0x1)), + _item(4, _func_types(func_ea)), + ] + mapping = {1: func_ea, 2: func_ea, 3: 0x1} + + failed = ImportDataTypes().execute( + FunctionDataTypesList.model_construct(items=items), matched_function_mapping=mapping + ) + + assert failed == {2, 3, 4} + printed = ida_typeinf.print_type(func_ea, ida_typeinf.PRTYPE_1LINE) + assert "count" in printed diff --git a/tests/idalib/variable_sync/test_read.py b/tests/idalib/variable_sync/test_read.py new file mode 100644 index 0000000..460cbf9 --- /dev/null +++ b/tests/idalib/variable_sync/test_read.py @@ -0,0 +1,49 @@ +from unittest.mock import MagicMock + +import pytest + +import ida_hexrays +import idaapi +import idc + +from reai_toolkit.app.services.variable_sync.variable_sync_service import ( + _read_decompiler_function, +) + +pytestmark = pytest.mark.idalib + + +@pytest.fixture +def deci(): + fake = MagicMock() + fake.decompiler_available = ida_hexrays.init_hexrays_plugin() + fake.art_lifter.lower_addr.side_effect = lambda addr: addr + fake.art_lifter.lift.side_effect = lambda func: func + return fake + + +def test_read_function_headless(loaded_binary, deci): + ea = idc.get_name_ea_simple("sub_401020") + assert ea != idaapi.BADADDR + + func = _read_decompiler_function(deci, ea) + + assert func is not None + assert func.header is not None + assert func.header.name == "sub_401020" + if deci.decompiler_available: + assert func.dec_obj is not None + + +def test_read_function_missing_returns_none(loaded_binary, deci): + assert _read_decompiler_function(deci, 0x1) is None + + +def test_read_function_without_decompiler(loaded_binary, deci): + deci.decompiler_available = False + ea = idc.get_name_ea_simple("sub_401020") + + func = _read_decompiler_function(deci, ea) + + assert func is not None + assert func.header is not None diff --git a/tests/ui/ida_ui_runner.py b/tests/ui/ida_ui_runner.py new file mode 100644 index 0000000..6f35809 --- /dev/null +++ b/tests/ui/ida_ui_runner.py @@ -0,0 +1,172 @@ +import json +import os +import sys +import time +import traceback + + +def _add_paths() -> None: + root = os.environ["REAI_UI_ROOT"] + pyver = f"python{sys.version_info.major}.{sys.version_info.minor}" + suffixes = [ + "vendor", + "vendor/site-packages", + "vendor/Lib/site-packages", + f"vendor/lib/{pyver}/site-packages", + f"vendor/{pyver}/site-packages", + ] + for base in (root, os.path.join(root, "reai_toolkit")): + for suffix in suffixes: + path = os.path.join(base, suffix) + if os.path.isdir(path) and path not in sys.path: + sys.path.insert(0, path) + if root not in sys.path: + sys.path.insert(0, root) + + +def _run(report: dict) -> None: + from types import SimpleNamespace + from unittest.mock import MagicMock + + import ida_hexrays + import ida_segment + import idautils + import idc + + from revengai import ( + FunctionArgument, + FunctionDataTypesList, + FunctionDataTypesListItem, + FunctionHeader, + FunctionInfo, + FunctionType, + Structure, + StructureMember, + ) + + from reai_toolkit.app.services.variable_sync.variable_sync_service import ( + _read_decompiler_function, + ) + from reai_toolkit.app.transformations.import_data_types import ImportDataTypes + + if not ida_hexrays.init_hexrays_plugin(): + report["errors"].append("hexrays unavailable") + return + + funcs = [ + ea + for ea in idautils.Functions() + if ida_segment.getseg(ea) and ida_segment.getseg(ea).type == ida_segment.SEG_CODE + ] + if len(funcs) < 3: + report["errors"].append(f"not enough code functions: {len(funcs)}") + return + + anchor, targets = funcs[0], funcs[1:41] + report["target_count"] = len(targets) + vdui = ida_hexrays.open_pseudocode(anchor, 0) + report["baseline_ea"] = vdui.cfunc.entry_ea + + class Recorder(ida_hexrays.Hexrays_Hooks): + def open_pseudocode(self, view) -> int: + report["pseudocode_opens"].append(view.cfunc.entry_ea) + return 0 + + def switch_pseudocode(self, view) -> int: + report["pseudocode_switches"].append(view.cfunc.entry_ea) + return 0 + + def refresh_pseudocode(self, view) -> int: + report["refresh_count"] += 1 + return 0 + + recorder = Recorder() + recorder.hook() + report["screen_ea_before"] = idc.get_screen_ea() + try: + struct_dep = Structure( + name="ReaiUiTestStruct", + size=8, + members={"0x0": StructureMember(name="field0", offset=0, type="int", size=4)}, + ) + items = [] + mapping: dict[int, int] = {} + for fid, ea in enumerate(targets, start=1): + args = { + "0x0": FunctionArgument(name="a", offset=0, size=4, type="int"), + "0x1": FunctionArgument(name="b", offset=1, size=8, type="ReaiUiTestStruct *"), + } + header = FunctionHeader(addr=ea, args=args, name=f"reai_ui_{fid}", type="int") + func_types = FunctionType( + addr=ea, header=header, name=f"reai_ui_{fid}", size=16, type="int" + ) + items.append( + FunctionDataTypesListItem.model_construct( + function_id=fid, + data_types=FunctionInfo.model_construct( + func_deps=[SimpleNamespace(actual_instance=struct_dep)], + func_types=func_types, + ), + ) + ) + mapping[fid] = ea + + started = time.monotonic() + failed = ImportDataTypes().execute( + FunctionDataTypesList.model_construct(items=items), + matched_function_mapping=mapping, + ) + report["import_seconds"] = round(time.monotonic() - started, 3) + report["import_failed_ids"] = sorted(failed) + report["functions_imported"] = len(targets) - len(failed) + + deci = MagicMock() + deci.decompiler_available = True + deci.art_lifter.lower_addr.side_effect = lambda addr: addr + deci.art_lifter.lift.side_effect = lambda func: func + started = time.monotonic() + report["functions_read"] = sum( + 1 for ea in targets if _read_decompiler_function(deci, ea) is not None + ) + report["read_seconds"] = round(time.monotonic() - started, 3) + finally: + recorder.unhook() + + report["final_vdui_ea"] = vdui.cfunc.entry_ea + report["screen_ea_after"] = idc.get_screen_ea() + report["ok"] = not report["errors"] + + +def main() -> None: + import ida_auto + import ida_pro + + report = { + "ok": False, + "errors": [], + "pseudocode_opens": [], + "pseudocode_switches": [], + "refresh_count": 0, + "baseline_ea": None, + "final_vdui_ea": None, + "screen_ea_before": None, + "screen_ea_after": None, + "target_count": 0, + "import_failed_ids": None, + "functions_imported": 0, + "functions_read": 0, + "import_seconds": None, + "read_seconds": None, + } + ida_auto.auto_wait() + try: + _add_paths() + _run(report) + except Exception: + report["errors"].append(traceback.format_exc()) + with open(os.environ["REAI_UI_REPORT"], "w") as fh: + json.dump(report, fh) + ida_pro.qexit(0) + + +main() diff --git a/tests/ui/test_gui_sync.py b/tests/ui/test_gui_sync.py new file mode 100644 index 0000000..f70e9bc --- /dev/null +++ b/tests/ui/test_gui_sync.py @@ -0,0 +1,93 @@ +import json +import os +import shutil +import subprocess +import sys + +import pytest + +pytestmark = pytest.mark.ida_ui + +ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +RUNNER = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ida_ui_runner.py") +HELLO_ELF = os.path.join(ROOT, "tests", "fixtures", "hello.elf") + + +def _ida_gui_binary() -> str | None: + path = os.path.join(os.environ.get("IDADIR", ""), "ida") + return path if os.path.isfile(path) else None + + +def _ida_gui_running() -> bool: + for name in ("ida", "ida64", "idat"): + try: + if subprocess.run(["pgrep", "-x", name], capture_output=True).returncode == 0: + return True + except FileNotFoundError: + return False + return False + + +def _headless_env() -> dict[str, str]: + if sys.platform == "linux" and not os.environ.get("DISPLAY") and not os.environ.get("WAYLAND_DISPLAY"): + return {"QT_QPA_PLATFORM": "offscreen"} + return {} + + +@pytest.mark.skipif( + os.environ.get("REAI_UI_TESTS") != "1", + reason="GUI IDA test; set REAI_UI_TESTS=1 to enable", +) +def test_sync_never_opens_or_switches_pseudocode(tmp_path): + ida = _ida_gui_binary() + if ida is None: + pytest.skip("IDADIR not set or ida binary missing") + if _ida_gui_running(): + pytest.skip("an IDA instance is already running (license seat busy)") + if not os.path.isfile(HELLO_ELF): + pytest.skip(f"missing fixture {HELLO_ELF}") + + binary = tmp_path / "hello.elf" + shutil.copy(HELLO_ELF, binary) + report_path = tmp_path / "report.json" + log_path = tmp_path / "ida.log" + env = dict( + os.environ, + REAI_UI_REPORT=str(report_path), + REAI_UI_ROOT=ROOT, + **_headless_env(), + ) + + proc = subprocess.run( + [ida, "-A", f"-S{RUNNER}", f"-L{log_path}", str(binary)], + env=env, + capture_output=True, + timeout=300, + ) + + if not report_path.is_file(): + log = log_path.read_text() if log_path.is_file() else "" + stderr = proc.stderr.decode(errors="replace") + stdout = proc.stdout.decode(errors="replace") + if "license" in log.lower() or "license" in stderr.lower(): + pytest.skip("IDA license unavailable") + pytest.fail( + f"IDA produced no report (rc={proc.returncode});" + f"\nstderr tail:\n{stderr[-2000:]}" + f"\nstdout tail:\n{stdout[-2000:]}" + f"\nlog tail:\n{log[-2000:]}" + ) + + report = json.loads(report_path.read_text()) + + assert report["errors"] == [] + assert report["ok"] is True + assert report["target_count"] > 0 + assert report["import_failed_ids"] == [] + assert report["functions_imported"] == report["target_count"] + assert report["functions_read"] == report["target_count"] + + assert report["pseudocode_opens"] == [] + assert report["pseudocode_switches"] == [] + assert report["final_vdui_ea"] == report["baseline_ea"] + assert report["screen_ea_after"] == report["screen_ea_before"] diff --git a/tests/unit/data_types/test_import_transform.py b/tests/unit/data_types/test_import_transform.py new file mode 100644 index 0000000..3253c5d --- /dev/null +++ b/tests/unit/data_types/test_import_transform.py @@ -0,0 +1,146 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from reai_toolkit.app.transformations import import_data_types as mod +from reai_toolkit.app.transformations.import_data_types import ( + APPLY_CHUNK_SIZE, + ImportDataTypes, +) +from revengai import ( + Enumeration, + FunctionDataTypesList, + FunctionInfo, + Structure, + StructureMember, + TypeDefinition, +) + + +@pytest.fixture +def deci(mocker): + instance = MagicMock() + mocker.patch.object(mod.DecompilerInterface, "discover", return_value=instance) + return instance + + +def _item(function_id: int, func_types=None, func_deps=None): + return SimpleNamespace( + function_id=function_id, + data_types=FunctionInfo.model_construct( + func_deps=[SimpleNamespace(actual_instance=d) for d in (func_deps or [])], + func_types=func_types, + ), + ) + + +def _functions(items): + return FunctionDataTypesList.model_construct(items=items) + + +def _struct(name: str, member_type: str = "int") -> Structure: + return Structure( + name=name, + size=8, + members={"0x0": StructureMember(name="field0", offset=0, type=member_type, size=8)}, + ) + + +def test_execute_no_items_skips_everything(deci): + idt = ImportDataTypes() + + assert idt.execute(_functions([SimpleNamespace(function_id=1, data_types=None)])) == set() + mod.DecompilerInterface.discover.assert_not_called() + + +def test_execute_skips_discover_without_dependencies(deci, mocker): + apply = mocker.patch.object(ImportDataTypes, "apply_function_type", return_value=True) + idt = ImportDataTypes() + + failed = idt.execute(_functions([_item(1, func_types=MagicMock(addr=0x1000))])) + + assert failed == set() + apply.assert_called_once() + mod.DecompilerInterface.discover.assert_not_called() + + +def test_execute_applies_shared_dependency_once(deci, mocker): + mocker.patch.object(ImportDataTypes, "apply_function_type", return_value=True) + shared = _struct("SharedStruct") + items = [ + _item(1, func_types=MagicMock(addr=0x1000), func_deps=[shared]), + _item(2, func_types=MagicMock(addr=0x2000), func_deps=[shared]), + ] + + ImportDataTypes().execute(_functions(items)) + + struct_writes = [c for c in deci.structs.mock_calls if "__setitem__" in str(c)] + assert len(struct_writes) == 1 + + +def test_execute_applies_subdependency_before_parent(deci, mocker): + mocker.patch.object(ImportDataTypes, "apply_function_type", return_value=True) + typedef = TypeDefinition(name="td_t", type="int") + parent = _struct("Parent", member_type="td_t") + enum = Enumeration(name="Colors", members={"RED": 0}) + items = [_item(1, func_types=MagicMock(addr=0x1000), func_deps=[parent, typedef, enum])] + + ImportDataTypes().execute(_functions(items)) + + writes = [c[0] for c in deci.mock_calls if "__setitem__" in c[0]] + assert writes.index("typedefs.__setitem__") < writes.index("structs.__setitem__") + assert "enums.__setitem__" in writes + + +def test_execute_survives_dependency_failure(deci, mocker): + apply = mocker.patch.object(ImportDataTypes, "apply_function_type", return_value=True) + deci.enums.__setitem__.side_effect = RuntimeError("til write failed") + enum = Enumeration(name="Colors", members={"RED": 0}) + items = [_item(1, func_types=MagicMock(addr=0x1000), func_deps=[enum])] + + failed = ImportDataTypes().execute(_functions(items)) + + assert failed == set() + apply.assert_called_once() + + +def test_execute_chunks_and_aggregates_failures(deci, mocker): + total = APPLY_CHUNK_SIZE + 5 + apply = mocker.patch.object( + ImportDataTypes, "apply_function_type", side_effect=lambda func, ea: ea % 2 == 0 + ) + items = [_item(fid, func_types=MagicMock(addr=fid)) for fid in range(total)] + + failed = ImportDataTypes().execute(_functions(items)) + + assert apply.call_count == total + assert failed == {fid for fid in range(total) if fid % 2 == 1} + + +def test_execute_uses_mapping_and_fails_unmapped(deci, mocker): + seen: list[int] = [] + + def record(func, ea): + seen.append(ea) + return True + + mocker.patch.object(ImportDataTypes, "apply_function_type", side_effect=record) + items = [ + _item(1, func_types=MagicMock(addr=0x1000)), + _item(2, func_types=MagicMock(addr=0x2000)), + ] + + failed = ImportDataTypes().execute(_functions(items), matched_function_mapping={1: 0x9000}) + + assert seen == [0x9000] + assert failed == {2} + + +def test_execute_skips_items_without_func_types(deci, mocker): + apply = mocker.patch.object(ImportDataTypes, "apply_function_type", return_value=True) + + failed = ImportDataTypes().execute(_functions([_item(1, func_types=None)])) + + assert failed == set() + apply.assert_not_called()