diff --git a/reai_toolkit/app/coordinators/chat_coordinator.py b/reai_toolkit/app/coordinators/chat_coordinator.py index 59e510c..63b9aa9 100644 --- a/reai_toolkit/app/coordinators/chat_coordinator.py +++ b/reai_toolkit/app/coordinators/chat_coordinator.py @@ -43,6 +43,8 @@ from reai_toolkit.app.coordinators.ai_decomp_coordinator import AiDecompCoordinator AI_DECOMP_TOOL_HINTS = ("decomp",) +EDIT_FUNCTION_TYPES_TOOL = "edit_function_types" +DATA_TYPES_SYNC_DEBOUNCE_S = 0.5 class ChatCoordinator(BaseCoordinator): @@ -64,6 +66,9 @@ def __init__( self._last_event_id: Optional[int] = None self._last_context: Optional[ConversationContextDTO] = None self._context_hooks: Optional[ChatContextHooks] = None + self._pending_type_ids: set[int] = set() + self._type_sync_timer: Optional[threading.Timer] = None + self._type_sync_lock = threading.Lock() @execute_ui def run_dialog(self, prefill_context: bool = False) -> None: @@ -100,6 +105,11 @@ def _wire_panel(self, panel: ChatPanel) -> None: def _on_pane_closed(self) -> None: if self._panel is not None: self._panel.stop_stream_worker() + with self._type_sync_lock: + if self._type_sync_timer is not None: + self._type_sync_timer.cancel() + self._type_sync_timer = None + self._pending_type_ids.clear() self._disable_context_tracking() self._panel = None self.log.info("Agent Chat panel closed.") @@ -226,12 +236,58 @@ def on_stream_finished(self) -> None: def _handle_tool_result(self, ev) -> None: if ev.updated: func_ids = [i for u in ev.updated if u.type == "function" for i in u.ids] - if func_ids: + if func_ids and self._is_edit_types_tool(ev): + self._sync_data_types(func_ids) + elif func_ids: self._sync_functions(func_ids) elif any(u.type == "analysis" for u in ev.updated): self.refresh_disassembly_view() self._maybe_open_viewer(ev) + @staticmethod + def _is_edit_types_tool(ev) -> bool: + return (ev.tool_name or "").lower() == EDIT_FUNCTION_TYPES_TOOL + + def _sync_data_types(self, func_ids: list) -> None: + with self._type_sync_lock: + self._pending_type_ids.update(int(f) for f in func_ids) + if self._type_sync_timer is not None: + self._type_sync_timer.cancel() + self._type_sync_timer = threading.Timer( + DATA_TYPES_SYNC_DEBOUNCE_S, self._flush_data_types + ) + self._type_sync_timer.daemon = True + self._type_sync_timer.start() + + def _flush_data_types(self) -> None: + with self._type_sync_lock: + ids = list(self._pending_type_ids) + self._pending_type_ids.clear() + self._type_sync_timer = None + if not ids: + return + func_map = self.app.netstore_service.get_function_mapping() + if func_map is None: + return + matches: dict[int, int] = {} + for fid in ids: + ea = func_map.function_map.get(str(fid)) + if ea is not None: + matches[fid] = ea + if not matches: + return + result = self.app.data_types_service.import_data_types( + matches, apply_stack_vars=True + ) + if result.error: + self.log.warning(f"Failed to sync data types: {result.error}") + return + self.refresh_disassembly_view() + last_ea = next(iter(matches.values())) + self._refresh_context_chip(last_ea) + if self._ai_decomp_coord is not None: + self._ai_decomp_coord.follow_function(last_ea) + def _sync_functions(self, func_ids: list) -> None: def _work() -> None: func_map = self.app.netstore_service.get_function_mapping() 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 d22b831..58159fe 100644 --- a/reai_toolkit/app/services/data_types/data_types_service.py +++ b/reai_toolkit/app/services/data_types/data_types_service.py @@ -39,7 +39,9 @@ def _import_worker(self, _: threading.Event, matches: dict[int, int]) -> None: if result.error: logger.error(f"RevEng.AI: {result.error}") - def import_data_types(self, matches: dict[int, int]) -> DataTypesImportResult: + def import_data_types( + self, matches: dict[int, int], apply_stack_vars: bool = False + ) -> DataTypesImportResult: if len(matches) == 0: return DataTypesImportResult() @@ -67,7 +69,9 @@ def import_data_types(self, matches: dict[int, int]) -> DataTypesImportResult: apply_failed_ids: set[int] = set() if response: - apply_failed_ids = idt.execute(response, matched_function_mapping=matches) or set() + apply_failed_ids = idt.execute( + response, matched_function_mapping=matches, apply_stack_vars=apply_stack_vars + ) or set() return DataTypesImportResult( remote_absent_ids=remote_absent_ids, diff --git a/reai_toolkit/app/transformations/import_data_types.py b/reai_toolkit/app/transformations/import_data_types.py index 59cd2b4..a7f4caa 100644 --- a/reai_toolkit/app/transformations/import_data_types.py +++ b/reai_toolkit/app/transformations/import_data_types.py @@ -1,3 +1,4 @@ +import re from typing import cast import ida_funcs @@ -19,6 +20,7 @@ ) APPLY_CHUNK_SIZE = 50 +_ANALYSIS_SCOPE_RE = re.compile(r"^[0-9a-fA-F]{64}(?:::|/)") class TaggedDependency: @@ -33,9 +35,15 @@ def __repr__(self) -> str: class ImportDataTypes: def __init__(self) -> None: - self.deci: DecompilerInterface - - def execute(self, functions: FunctionDataTypesList, matched_function_mapping: dict[int, int] = {}) -> set[int]: + self.deci: DecompilerInterface | None = None + self._stack_vars_ok: bool | None = None + + def execute( + self, + functions: FunctionDataTypesList, + matched_function_mapping: dict[int, int] = {}, + apply_stack_vars: bool = False, + ) -> set[int]: items: list[FunctionDataTypesListItem] = [ item for item in functions.items if item.data_types is not None ] @@ -54,7 +62,7 @@ def execute(self, functions: FunctionDataTypesList, matched_function_mapping: di 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) + failed |= self._apply_chunk(chunk, matched_function_mapping, apply_stack_vars) logger.info( f"RevEng.AI: applied data types to {min(start + APPLY_CHUNK_SIZE, total)}/{total} functions" ) @@ -73,9 +81,33 @@ def _build_lookup(self, items: list[FunctionDataTypesListItem]) -> dict[str, Tag return lookup + def _ensure_deci(self) -> None: + if self.deci is None: + self.deci = DecompilerInterface.discover(force_decompiler="ida") # type: ignore + + def _stack_vars_available(self, ea: int) -> bool: + if self._stack_vars_ok is None: + self._stack_vars_ok = self._probe_decompiler(ea) + if not self._stack_vars_ok: + logger.info( + "RevEng.AI: decompiler unavailable for this binary; skipping stack variable sync" + ) + return self._stack_vars_ok + + @staticmethod + def _probe_decompiler(ea: int) -> bool: + try: + import ida_hexrays + + if not ida_hexrays.init_hexrays_plugin(): + return False + return ida_hexrays.decompile(ea, ida_hexrays.hexrays_failure_t()) is not None + except Exception: + return False + @execute_write def _apply_dependencies(self, lookup: dict[str, TaggedDependency]) -> None: - self.deci = DecompilerInterface.discover(force_decompiler="ida") # type: ignore + self._ensure_deci() for tagged_dependency in lookup.values(): try: self.process_dependency(tagged_dependency, lookup) @@ -87,7 +119,10 @@ def _apply_dependencies(self, lookup: dict[str, TaggedDependency]) -> None: @execute_write def _apply_chunk( - self, chunk: list[FunctionDataTypesListItem], matched_function_mapping: dict[int, int] + self, + chunk: list[FunctionDataTypesListItem], + matched_function_mapping: dict[int, int], + apply_stack_vars: bool = False, ) -> set[int]: failed: set[int] = set() for item in chunk: @@ -103,6 +138,8 @@ def _apply_chunk( if not self.apply_function_type(func, ea): failed.add(item.function_id) + if apply_stack_vars: + self.apply_stack_variables(func, ea) except Exception as e: logger.warning( f"RevEng.AI: skipped data types for function {item.function_id}: {e!r}" @@ -149,6 +186,38 @@ def apply_function_type(self, func: FunctionType, ea: int) -> bool: return bool(ida_typeinf.apply_tinfo(ea, proto, ida_typeinf.TINFO_DEFINITE)) + def apply_stack_variables(self, func: FunctionType, ea: int) -> None: + stack_vars = getattr(func, "stack_vars", None) + if not isinstance(stack_vars, dict) or not stack_vars: + return + + if not self._stack_vars_available(ea): + return + + self._ensure_deci() + if self.deci is None: + return + + try: + lifted_ea: int = self.deci.art_lifter.lift_addr(ea) + svars: dict[int, libbs.artifacts.StackVariable] = { + svar.offset: libbs.artifacts.StackVariable( + stack_offset=svar.offset, + name=svar.name, + type_=self.normalise_type(svar.type) if svar.type else None, + size=svar.size, + addr=lifted_ea, + ) + for svar in stack_vars.values() + } + self.deci.functions[lifted_ea] = libbs.artifacts.Function( + addr=lifted_ea, + header=libbs.artifacts.FunctionHeader(addr=lifted_ea), + stack_vars=svars, + ) + except Exception as e: + logger.warning(f"RevEng.AI: skipped stack variables for 0x{ea:x}: {e!r}") + @staticmethod def _current_func_details(ea: int) -> "ida_typeinf.func_type_data_t": details = ida_typeinf.func_type_data_t() @@ -223,6 +292,8 @@ def normalise_type(data_type: str) -> str: pos: int = data_type.find(delimiter) data_type = data_type[pos+len(delimiter):] + data_type = _ANALYSIS_SCOPE_RE.sub("", data_type) + # TODO: PLU-213 Add IDA typedefs for Ghidra primitives so we don't need to bother doing this... if data_type == "uchar": data_type = "unsigned char" diff --git a/tests/unit/data_types/test_import_transform.py b/tests/unit/data_types/test_import_transform.py index 3253c5d..0ed0362 100644 --- a/tests/unit/data_types/test_import_transform.py +++ b/tests/unit/data_types/test_import_transform.py @@ -21,6 +21,7 @@ @pytest.fixture def deci(mocker): instance = MagicMock() + instance.art_lifter.lift_addr.side_effect = lambda addr: addr mocker.patch.object(mod.DecompilerInterface, "discover", return_value=instance) return instance @@ -144,3 +145,75 @@ def test_execute_skips_items_without_func_types(deci, mocker): assert failed == set() apply.assert_not_called() + + +_HASH = "259156281adba01eb86070f77a039e7054f268c973326adcee5fe4533f14b292" + + +@pytest.mark.parametrize( + "raw,expected", + [ + (f"{_HASH}::Candidate *", "Candidate *"), + (f"{_HASH}/std::vector_>", "std::vector_>"), + (f"{_HASH}::_Tree_node::_Node *", "_Tree_node::_Node *"), + ("DWARF/stdint.h::uint32_t", "uint32_t"), + ("std::vector", "std::vector"), + ("int", "int"), + ], +) +def test_normalise_type_strips_analysis_scope(raw, expected): + assert ImportDataTypes.normalise_type(raw) == expected + + +def _svar(offset: int, name: str, type_str: str, size: int = 4): + return SimpleNamespace(offset=offset, name=name, type=type_str, size=size) + + +def test_apply_stack_variables_writes_function_with_normalised_types(deci, mocker): + mocker.patch.object(ImportDataTypes, "_probe_decompiler", return_value=True) + func = SimpleNamespace( + stack_vars={ + "0x4": _svar(4, "lhs", "int"), + "0x8": _svar(8, "rhs", f"{_HASH}::Candidate *"), + } + ) + + ImportDataTypes().apply_stack_variables(func, 0x1000) + + deci.functions.__setitem__.assert_called_once() + ea, written = deci.functions.__setitem__.call_args.args + assert ea == 0x1000 + assert set(written.stack_vars) == {4, 8} + assert written.stack_vars[4].name == "lhs" + assert written.stack_vars[8].type == "Candidate *" + + +def test_apply_stack_variables_noop_without_stack_vars(deci): + ImportDataTypes().apply_stack_variables(SimpleNamespace(stack_vars=None), 0x1000) + ImportDataTypes().apply_stack_variables(SimpleNamespace(stack_vars={}), 0x1000) + + mod.DecompilerInterface.discover.assert_not_called() + deci.functions.__setitem__.assert_not_called() + + +def test_execute_applies_stack_vars_only_when_enabled(deci, mocker): + mocker.patch.object(ImportDataTypes, "apply_function_type", return_value=True) + svapply = mocker.patch.object(ImportDataTypes, "apply_stack_variables") + items = [_item(1, func_types=MagicMock(addr=0x1000))] + + ImportDataTypes().execute(_functions(items)) + svapply.assert_not_called() + + ImportDataTypes().execute(_functions(items), apply_stack_vars=True) + svapply.assert_called_once() + + +def test_apply_stack_variables_skips_when_decompiler_unavailable(deci, mocker): + mocker.patch.object(ImportDataTypes, "_probe_decompiler", return_value=False) + + ImportDataTypes().apply_stack_variables( + SimpleNamespace(stack_vars={"0x4": _svar(4, "lhs", "int")}), 0x1000 + ) + + mod.DecompilerInterface.discover.assert_not_called() + deci.functions.__setitem__.assert_not_called()