From 8f3cccfb352b0253bb1585cdbf3d5a1aa50ec4e9 Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Wed, 9 Sep 2026 11:32:53 +0100 Subject: [PATCH 01/12] feat(PLU-305): sdk version bump --- pyproject.toml | 2 +- uv.lock | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 27ce0c6..875799e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ dependencies = [ "requests>=2.32", "loguru>=0.7", "pydantic", - "revengai>=3.123.0", + "revengai>=4.20.0", "libbs>=2.16.5", ] diff --git a/uv.lock b/uv.lock index 83a1f97..daa5301 100644 --- a/uv.lock +++ b/uv.lock @@ -820,7 +820,7 @@ requires-dist = [ { name = "loguru", specifier = ">=0.7" }, { name = "pydantic" }, { name = "requests", specifier = ">=2.32" }, - { name = "revengai", specifier = ">=3.123.0" }, + { name = "revengai", specifier = ">=4.20.0" }, ] [package.metadata.requires-dev] @@ -855,7 +855,7 @@ wheels = [ [[package]] name = "revengai" -version = "3.123.0" +version = "4.20.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "lazy-imports" }, @@ -864,9 +864,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e4/11/c263645fd280bb0d4c0a3cb3057880570bd8d6581b71be0f688f7adc6dee/revengai-3.123.0.tar.gz", hash = "sha256:871199b75675ac54543005179e7f92b845cfceba79ee49c8f11db0d77613aa7f", size = 369727, upload-time = "2026-07-23T08:49:18.527Z" } +sdist = { url = "https://files.pythonhosted.org/packages/52/29/9617fe378a9d191d84c62d1f4ad96411dc55ad516c5470b5d95ef27dc5c3/revengai-4.20.0.tar.gz", hash = "sha256:84886a55baf878c7a31f38eeb7199e8b29fa10077378c33c30bc84d64aa646c3", size = 426999, upload-time = "2026-09-02T14:37:02.389Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/91/e1/5a01283b4eb34f54879ce7a91a2db81d27f987c0c1a5e33ce2e81c6ae297/revengai-3.123.0-py3-none-any.whl", hash = "sha256:f488e653bd88b9c8d432d062d11417710502b9bb10966c88c2c77f69e588fdb1", size = 966561, upload-time = "2026-07-23T08:49:17.023Z" }, + { url = "https://files.pythonhosted.org/packages/47/ef/5731372d613239278216f5c152060cfb9091cdc0e5ddebd6d2ee5cf9475d/revengai-4.20.0-py3-none-any.whl", hash = "sha256:a6c362fb691748e1fa8b82a82d232c43a703c4112808a8ab2a98b93b3d24898f", size = 1098020, upload-time = "2026-09-02T14:37:01.028Z" }, ] [[package]] From a7059df36836d83d403e82c0e90591a745a72f01 Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Wed, 9 Sep 2026 12:04:18 +0100 Subject: [PATCH 02/12] feat: data types pull path refactor --- .../app/coordinators/chat_coordinator.py | 4 +- .../services/data_types/data_types_service.py | 47 +- .../app/transformations/import_data_types.py | 452 ++++++++++-------- tests/idalib/data_types/test_import.py | 125 +++-- tests/ui/ida_ui_runner.py | 70 +-- .../unit/data_types/test_import_transform.py | 261 ++++++---- tests/unit/data_types/test_sdk_schemas.py | 93 ++-- tests/unit/data_types/test_service.py | 108 +++-- 8 files changed, 671 insertions(+), 489 deletions(-) diff --git a/reai_toolkit/app/coordinators/chat_coordinator.py b/reai_toolkit/app/coordinators/chat_coordinator.py index 63b9aa9..8284ca2 100644 --- a/reai_toolkit/app/coordinators/chat_coordinator.py +++ b/reai_toolkit/app/coordinators/chat_coordinator.py @@ -276,9 +276,7 @@ def _flush_data_types(self) -> None: matches[fid] = ea if not matches: return - result = self.app.data_types_service.import_data_types( - matches, apply_stack_vars=True - ) + result = self.app.data_types_service.import_data_types(matches) if result.error: self.log.warning(f"Failed to sync data types: {result.error}") return 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 58159fe..cd0751b 100644 --- a/reai_toolkit/app/services/data_types/data_types_service.py +++ b/reai_toolkit/app/services/data_types/data_types_service.py @@ -7,14 +7,13 @@ from loguru import logger from reai_toolkit.app.core.netstore_service import SimpleNetStore -from revengai import ( - FunctionsDataTypesApi, - FunctionDataTypesList, - BaseResponseFunctionDataTypesList -) +from revengai import DataTypesApi from reai_toolkit.app.interfaces.thread_service import IThreadService -from reai_toolkit.app.transformations.import_data_types import ImportDataTypes +from reai_toolkit.app.transformations.import_data_types import ( + FunctionSignatures, + ImportDataTypes, +) FUNCTION_IDS_BATCH_SIZE = 50 @@ -39,9 +38,7 @@ 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], apply_stack_vars: bool = False - ) -> DataTypesImportResult: + def import_data_types(self, matches: dict[int, int]) -> DataTypesImportResult: if len(matches) == 0: return DataTypesImportResult() @@ -49,7 +46,7 @@ def import_data_types( matched_function_ids: list[int] = list(matches.keys()) try: - response: FunctionDataTypesList | None = self._get_data_types(matched_function_ids) + response: FunctionSignatures | None = self._get_data_types(matched_function_ids) except NotFoundException as e: logger.warning(f"failed to apply data types for {len(matched_function_ids)} functions: {e}") return DataTypesImportResult(remote_absent_ids=set(matched_function_ids)) @@ -61,7 +58,7 @@ def import_data_types( return DataTypesImportResult(error=f"Failed to sync function data types: {e}") present_ids: set[int] = ( - {item.function_id for item in response.items if item.data_types is not None} + {item["function_id"] for item in response.items if item.get("has_signature")} if response else set() ) @@ -70,7 +67,7 @@ def import_data_types( apply_failed_ids: set[int] = set() if response: apply_failed_ids = idt.execute( - response, matched_function_mapping=matches, apply_stack_vars=apply_stack_vars + response, matched_function_mapping=matches ) or set() return DataTypesImportResult( @@ -78,19 +75,25 @@ def import_data_types( apply_failed_ids=apply_failed_ids, ) - def _get_data_types(self, function_ids: list[int] | None = None) -> FunctionDataTypesList | None: + def _get_data_types(self, function_ids: list[int] | None = None) -> FunctionSignatures | None: if not function_ids: return None - items = [] + signatures = FunctionSignatures() with self.yield_api_client(sdk_config=self.sdk_config) as api_client: - client = FunctionsDataTypesApi(api_client=api_client) + client = DataTypesApi(api_client=api_client) for start in range(0, len(function_ids), FUNCTION_IDS_BATCH_SIZE): chunk = function_ids[start:start + FUNCTION_IDS_BATCH_SIZE] - response: BaseResponseFunctionDataTypesList = ( - client.list_function_data_types_for_functions(function_ids=chunk) # type: ignore - ) - if response.status and response.data: - items.extend(response.data.items) - - return FunctionDataTypesList(items=items) + response = client.v3_list_function_signatures( + function_ids=chunk, include_data_types=True + ).to_dict() + + signatures.items.extend(response.get("items") or []) + for group in response.get("data_types") or []: + for entry in group.get("items") or []: + data_type_id = entry.get("data_type_id") + if data_type_id is None: + continue + signatures.data_types[str(data_type_id)] = entry + + return signatures diff --git a/reai_toolkit/app/transformations/import_data_types.py b/reai_toolkit/app/transformations/import_data_types.py index a7f4caa..e90c41d 100644 --- a/reai_toolkit/app/transformations/import_data_types.py +++ b/reai_toolkit/app/transformations/import_data_types.py @@ -1,5 +1,6 @@ import re -from typing import cast +from dataclasses import dataclass, field +from typing import Any, Optional import ida_funcs import ida_typeinf @@ -9,163 +10,322 @@ from libbs.api import DecompilerInterface from libbs.decompilers.ida.compat import execute_write, convert_type_str_to_ida_type from loguru import logger -from revengai import ( - Enumeration, - FunctionArgument, - FunctionDataTypesList, - FunctionDataTypesListItem, - FunctionType, - Structure, - TypeDefinition, -) APPLY_CHUNK_SIZE = 50 _ANALYSIS_SCOPE_RE = re.compile(r"^[0-9a-fA-F]{64}(?:::|/)") +_RECORD_KINDS = frozenset({"STRUCT", "UNION"}) + +_GHIDRA_PRIMITIVE_DECLARATORS: tuple[str, ...] = ( + "unsigned __int8 byte", + "__int8 sbyte", + "unsigned char uchar", + "signed char schar", + "unsigned __int16 word", + "__int16 sword", + "unsigned __int32 dword", + "__int32 sdword", + "unsigned __int64 qword", + "__int64 sqword", + "__int64 longlong", + "unsigned __int64 ulonglong", + "unsigned __int8 undefined", + "unsigned __int8 undefined1", + "unsigned __int16 undefined2", + "unsigned __int32 undefined4", + "unsigned __int8 undefined6[6]", + "unsigned __int64 undefined8", + "unsigned __int32 pointer32", + "unsigned __int64 pointer64", + "unsigned __int16 wchar16", + "unsigned __int32 wchar32", + "_TBYTE float10", + "unsigned __int8 code", +) + +_GHIDRA_PRIMITIVE_NAME_RE = re.compile(r"(\w+)(?:\[\d+\])?$") + + +def _primitive_name(declarator: str) -> str: + match = _GHIDRA_PRIMITIVE_NAME_RE.search(declarator) + return match.group(1) if match else declarator + + +def install_ghidra_primitives() -> int: + missing = [ + f"typedef {declarator};" + for declarator in _GHIDRA_PRIMITIVE_DECLARATORS + if convert_type_str_to_ida_type(_primitive_name(declarator)) is None + ] + if not missing: + return 0 + + errors = ida_typeinf.parse_decls(None, "".join(missing), None, 0) + if errors: + logger.warning( + f"RevEng.AI: {errors} error(s) declaring Ghidra primitive types; " + "some server types may not resolve" + ) + return len(missing) + + +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` + # Let's remove the DWARF/*.h prefix + if data_type.startswith("DWARF/"): + # Find the first occurence of `::` + delimiter: str = "::" + pos: int = data_type.find(delimiter) + data_type = data_type[pos+len(delimiter):] + + return _ANALYSIS_SCOPE_RE.sub("", data_type) + + +@dataclass +class FunctionSignatures: + items: list[dict] = field(default_factory=list) + data_types: dict[str, dict] = field(default_factory=dict) + + +@dataclass +class _Prototype: + function_id: int + name: Optional[str] + return_type: Optional[str] + args: list[libbs.artifacts.FunctionArgument] + + +class _TypeResolver: + def __init__(self, data_types: dict[str, dict]) -> None: + self._data_types = data_types + self._artifacts: dict[str, Any] = {} + self._visiting: set[str] = set() + + @property + def artifacts(self) -> dict[str, Any]: + return self._artifacts + + def resolve(self, data_type_id: Optional[int]) -> tuple[Optional[str], Optional[int]]: + if data_type_id is None: + return None, None + + entry = self._data_types.get(str(data_type_id)) + if entry is None: + return None, None + + kind = (entry.get("kind") or "").upper() + definition = entry.get("definition") or {} + + if kind in _RECORD_KINDS: + self._record(entry, definition) + elif kind == "ENUM": + self._enum(entry, definition) + elif kind == "TYPEDEF": + self._typedef(entry, definition) + elif kind == "POINTER": + self.resolve(definition.get("pointee_data_type_id")) + elif kind == "ARRAY": + self.resolve(definition.get("element_data_type_id")) -class TaggedDependency: - def __init__(self, dependency: Structure | Enumeration | TypeDefinition) -> None: - self.dependency: Structure | Enumeration | TypeDefinition = dependency - self.processed: bool = False - self.name: str = self.dependency.name + name = entry.get("name") + return (normalise_type(name) if name else None), entry.get("size") - def __repr__(self) -> str: - return self.dependency.__repr__() + def _record(self, entry: dict, definition: dict) -> None: + name = entry.get("name") + if not name or name in self._artifacts or name in self._visiting: + return + + self._visiting.add(name) + members: dict[int, libbs.artifacts.StructMember] = {} + for member in definition.get("members") or []: + member_type, member_size = self.resolve(member.get("data_type_id")) + offset = member.get("offset") + if offset is None or member_type is None: + continue + members[offset] = libbs.artifacts.StructMember( + name=member.get("name") or f"field_{offset:x}", + offset=offset, + type_=member_type, + size=member.get("size") or member_size, + ) + self._visiting.discard(name) + + size = entry.get("size") + if size is None: + return + self._artifacts[name] = libbs.artifacts.Struct( + name=name, size=size, members=members + ) + + def _enum(self, entry: dict, definition: dict) -> None: + name = entry.get("name") + if not name or name in self._artifacts: + return + + members: dict[str, int] = {} + for value in definition.get("values") or []: + try: + members[value["name"]] = int(value["value"]) + except (KeyError, TypeError, ValueError): + continue + + self._artifacts[name] = libbs.artifacts.Enum(name=name, members=members) + + def _typedef(self, entry: dict, definition: dict) -> None: + name = entry.get("name") + if not name or name in self._artifacts or name in self._visiting: + return + + self._visiting.add(name) + target, _ = self.resolve(definition.get("target_data_type_id")) + self._visiting.discard(name) + + if target is None: + return + self._artifacts[name] = libbs.artifacts.Typedef(name=name, type_=target) class ImportDataTypes: def __init__(self) -> None: self.deci: DecompilerInterface | None = None - self._stack_vars_ok: bool | None = None def execute( self, - functions: FunctionDataTypesList, + signatures: FunctionSignatures, 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 + resolver = _TypeResolver(signatures.data_types) + prototypes = [ + proto + for proto in (self._prototype(item, resolver) for item in signatures.items) + if proto is not None ] - if not items: + if not prototypes: 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) + self._install_primitives() + + if resolver.artifacts: + self._apply_dependencies(resolver.artifacts) failed: set[int] = set() - total: int = len(items) + total: int = len(prototypes) 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, apply_stack_vars) + chunk = prototypes[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" ) return failed - 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 + @staticmethod + @execute_write + def _install_primitives() -> None: + declared = install_ghidra_primitives() + if declared: + logger.info(f"RevEng.AI: declared {declared} Ghidra primitive type(s)") - if dep.actual_instance.name not in lookup: - lookup[dep.actual_instance.name] = TaggedDependency(dep.actual_instance) # type: ignore + @staticmethod + def _prototype(item: dict, resolver: _TypeResolver) -> Optional[_Prototype]: + if not item.get("has_signature"): + return None + + function_id = item.get("function_id") + if function_id is None: + return None + + args: list[libbs.artifacts.FunctionArgument] = [] + for parameter in item.get("parameters") or []: + ordinal = parameter.get("ordinal") + if ordinal is None: + continue + arg_type, arg_size = resolver.resolve(parameter.get("data_type_id")) + args.append( + libbs.artifacts.FunctionArgument( + offset=ordinal, + name=parameter.get("name") or f"a{ordinal + 1}", + type_=arg_type, + size=arg_size, + ) + ) + + return_type, _ = resolver.resolve(item.get("return_data_type_id")) - return lookup + return _Prototype( + function_id=function_id, + name=item.get("function_name"), + return_type=return_type, + args=sorted(args, key=lambda a: a.offset), + ) 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: + def _apply_dependencies(self, artifacts: dict[str, Any]) -> None: self._ensure_deci() - for tagged_dependency in lookup.values(): + for name, artifact in artifacts.items(): try: - self.process_dependency(tagged_dependency, lookup) + self.apply_dependency(artifact) except Exception as e: - logger.warning( - f"RevEng.AI: skipped dependency {tagged_dependency.name!r}: {e!r}" - ) - tagged_dependency.processed = True + logger.warning(f"RevEng.AI: skipped dependency {name!r}: {e!r}") + + def apply_dependency(self, artifact: Any) -> None: + if self.deci is None: + return + + match artifact: + case libbs.artifacts.Struct(): + self.deci.structs[artifact.name] = artifact + case libbs.artifacts.Enum(): + self.deci.enums[artifact.name] = artifact + case libbs.artifacts.Typedef(): + self.deci.typedefs[artifact.name] = artifact + case _: + logger.warning(f"unsupported dependency type: {artifact}") @execute_write def _apply_chunk( self, - chunk: list[FunctionDataTypesListItem], + chunk: list[_Prototype], matched_function_mapping: dict[int, int], - apply_stack_vars: bool = False, ) -> set[int]: failed: set[int] = set() - for item in chunk: - func: FunctionType | None = item.data_types.func_types - if func is None: + for proto in chunk: + ea = matched_function_mapping.get(proto.function_id) + if ea is None: + failed.add(proto.function_id) continue 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) - if apply_stack_vars: - self.apply_stack_variables(func, ea) + if not self.apply_function_type(proto, ea): + failed.add(proto.function_id) except Exception as e: logger.warning( - f"RevEng.AI: skipped data types for function {item.function_id}: {e!r}" + f"RevEng.AI: skipped data types for function {proto.function_id}: {e!r}" ) - failed.add(item.function_id) + failed.add(proto.function_id) return failed - def apply_function_type(self, func: FunctionType, ea: int) -> bool: + def apply_function_type(self, proto: _Prototype, ea: int) -> bool: if ida_funcs.get_func(ea) is None: - logger.warning(f"failed to update function: {func.name} at 0x{ea:x}") + logger.warning(f"failed to update function: {proto.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 + for arg in proto.args: + arg_tif = convert_type_str_to_ida_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 proto.return_type: + ret_tif = convert_type_str_to_ida_type(proto.return_type) if ret_tif is None: return False details.rettype = ret_tif @@ -180,43 +340,11 @@ def apply_function_type(self, func: FunctionType, ea: int) -> bool: funcarg.type = arg_tif details.push_back(funcarg) - proto = ida_typeinf.tinfo_t() - if not proto.create_func(details): + prototype = ida_typeinf.tinfo_t() + if not prototype.create_func(details): return False - 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}") + return bool(ida_typeinf.apply_tinfo(ea, prototype, ida_typeinf.TINFO_DEFINITE)) @staticmethod def _current_func_details(ea: int) -> "ida_typeinf.func_type_data_t": @@ -235,71 +363,3 @@ def _current_func_details(ea: int) -> "ida_typeinf.func_type_data_t": 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(): - self.update_struct(cast(Structure, dependency), lookup) - case Enumeration(): - self.update_enum(cast(Enumeration, dependency)) - case TypeDefinition(): - self.update_typedef(cast(TypeDefinition, dependency), lookup) - case _: - logger.warning(f"unsupported dependency type: {dependency}") - - tagged_dependency.processed = True - - 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: - self.process_dependency(subdependency, lookup) - member.type = self.normalise_type(member.type) - - self.deci.structs[imported_struct.name] = libbs.artifacts.Struct( - name=imported_struct.name, size=imported_struct.size, members={v.offset: v for v in imported_struct.members.values()} # type: ignore - ) - - def update_enum(self, imported_enum: Enumeration) -> None: - self.deci.enums[imported_enum.name] = libbs.artifacts.Enum(name=imported_enum.name, members=imported_enum.members) - - def update_typedef(self, imported_typedef: TypeDefinition, lookup: dict[str, TaggedDependency]) -> None: - subdependency: TaggedDependency | None = lookup.get(imported_typedef.type) - if subdependency: - self.process_dependency(subdependency, lookup) - - normalized_type: str = self.normalise_type(imported_typedef.type) - self.deci.typedefs[imported_typedef.name] = libbs.artifacts.Typedef( - name=imported_typedef.name, type_=normalized_type - ) - - @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` - # Let's remove the DWARF/*.h prefix - if data_type.startswith("DWARF/"): - # Find the first occurence of `::` - delimiter: 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" - elif data_type == "qword": - data_type = "unsigned __int64" - elif data_type == "sqword": - data_type = "__int64" - - return data_type diff --git a/tests/idalib/data_types/test_import.py b/tests/idalib/data_types/test_import.py index f8c7ca4..b43a6d5 100644 --- a/tests/idalib/data_types/test_import.py +++ b/tests/idalib/data_types/test_import.py @@ -3,16 +3,14 @@ import ida_typeinf import idaapi import idc -from revengai import ( - FunctionArgument, - FunctionDataTypesList, - FunctionDataTypesListItem, - FunctionHeader, - FunctionInfo, - FunctionType, -) +import libbs.artifacts -from reai_toolkit.app.transformations.import_data_types import ImportDataTypes +from reai_toolkit.app.transformations.import_data_types import ( + FunctionSignatures, + ImportDataTypes, + _Prototype, + install_ghidra_primitives, +) pytestmark = pytest.mark.idalib @@ -30,37 +28,40 @@ def func_ea(loaded_binary): 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 _arg(offset: int, name: str, type_str: str) -> libbs.artifacts.FunctionArgument: + return libbs.artifacts.FunctionArgument( + offset=offset, name=name, type_=type_str, size=8 + ) -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, - ), +def _proto(function_id: int = 1, ret: str = "int", args: tuple = ()) -> _Prototype: + return _Prototype( + function_id=function_id, name="reai_test_func", - size=16, - type=ret, + return_type=ret, + args=list(args), ) -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 _signature(function_id: int, parameters, return_id=None): + return { + "function_id": function_id, + "function_name": "reai_test_func", + "has_signature": True, + "parameters": parameters, + "return_data_type_id": return_id, + } + + +def _base(data_type_id: int, name: str, size: int = 4): + return {"data_type_id": data_type_id, "name": name, "kind": "BASE", "size": size} 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 *"))) + proto = _proto(ret="int", args=(_arg(0, "count", "int"), _arg(1, "buf", "char *"))) - assert idt.apply_function_type(func, func_ea) is True + assert idt.apply_function_type(proto, func_ea) is True printed = ida_typeinf.print_type(func_ea, ida_typeinf.PRTYPE_1LINE) assert "count" in printed @@ -68,24 +69,49 @@ def test_apply_function_type_sets_prototype_with_named_args(func_ea): assert "char *" in printed -def test_apply_function_type_normalises_dwarf_types(func_ea): +def test_ghidra_primitive_names_resolve_after_install(loaded_binary): + install_ghidra_primitives() + + from libbs.decompilers.ida.compat import convert_type_str_to_ida_type + + for name, size in ( + ("byte", 1), + ("word", 2), + ("dword", 4), + ("qword", 8), + ("sqword", 8), + ("undefined", 1), + ("undefined6", 6), + ("float10", 10), + ): + tif = convert_type_str_to_ida_type(name) + assert tif is not None, f"{name} did not resolve" + assert tif.get_size() == size, f"{name} resolved at {tif.get_size()}, want {size}" + + +def test_install_ghidra_primitives_is_idempotent(loaded_binary): + install_ghidra_primitives() + + assert install_ghidra_primitives() == 0 + + +def test_apply_function_type_accepts_ghidra_primitives(func_ea): + install_ghidra_primitives() idt = ImportDataTypes() - func = _func_types(func_ea, ret="uchar", args=(_arg(0, "n", "DWARF/stdint.h::qword"),)) + proto = _proto(ret="uchar", args=(_arg(0, "n", "qword"), _arg(1, "flags", "dword"))) - assert idt.apply_function_type(func, func_ea) is True + assert idt.apply_function_type(proto, 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 + assert "n" in printed + assert "flags" 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(_proto(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 + assert idt.apply_function_type(_proto(ret="void", args=()), func_ea) is True printed = ida_typeinf.print_type(func_ea, ida_typeinf.PRTYPE_1LINE) assert "void" in printed @@ -94,29 +120,32 @@ def test_apply_function_type_keeps_args_when_remote_has_none(func_ea): def test_apply_function_type_rejects_unparseable_arg(func_ea): idt = ImportDataTypes() - func = _func_types(func_ea, args=(_arg(0, "x", "totally bogus type!!"),)) + proto = _proto(args=(_arg(0, "x", "totally bogus type!!"),)) - assert idt.apply_function_type(func, func_ea) is False + assert idt.apply_function_type(proto, 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 + assert idt.apply_function_type(_proto(), 0x1) is False def test_execute_applies_via_mapping_and_reports_failures(func_ea): + data_types = [_base(1, "int"), _base(2, "totally bogus type!!")] 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)), + _signature(1, [{"ordinal": 0, "name": "count", "data_type_id": 1}], return_id=1), + _signature(2, [{"ordinal": 0, "name": "x", "data_type_id": 2}], return_id=1), + _signature(3, [], return_id=1), + _signature(4, [], return_id=1), ] + signatures = FunctionSignatures( + items=items, + data_types={str(e["data_type_id"]): e for e in data_types}, + ) mapping = {1: func_ea, 2: func_ea, 3: 0x1} - failed = ImportDataTypes().execute( - FunctionDataTypesList.model_construct(items=items), matched_function_mapping=mapping - ) + failed = ImportDataTypes().execute(signatures, matched_function_mapping=mapping) assert failed == {2, 3, 4} printed = ida_typeinf.print_type(func_ea, ida_typeinf.PRTYPE_1LINE) diff --git a/tests/ui/ida_ui_runner.py b/tests/ui/ida_ui_runner.py index 6f35809..035345c 100644 --- a/tests/ui/ida_ui_runner.py +++ b/tests/ui/ida_ui_runner.py @@ -25,7 +25,6 @@ def _add_paths() -> None: def _run(report: dict) -> None: - from types import SimpleNamespace from unittest.mock import MagicMock import ida_hexrays @@ -33,21 +32,13 @@ def _run(report: dict) -> None: 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 + from reai_toolkit.app.transformations.import_data_types import ( + FunctionSignatures, + ImportDataTypes, + ) if not ida_hexrays.init_hexrays_plugin(): report["errors"].append("hexrays unavailable") @@ -84,36 +75,47 @@ def refresh_pseudocode(self, view) -> int: 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)}, - ) + data_types = { + "1": {"data_type_id": 1, "name": "int", "kind": "BASE", "size": 4}, + "2": { + "data_type_id": 2, + "name": "ReaiUiTestStruct", + "kind": "STRUCT", + "size": 8, + "definition": { + "members": [ + {"name": "field0", "offset": 0, "size": 4, "data_type_id": 1} + ] + }, + }, + "3": { + "data_type_id": 3, + "name": "ReaiUiTestStruct *", + "kind": "POINTER", + "size": 8, + "definition": {"pointee_data_type_id": 2}, + }, + } 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, - ), - ) + { + "function_id": fid, + "function_name": f"reai_ui_{fid}", + "has_signature": True, + "parameters": [ + {"ordinal": 0, "name": "a", "data_type_id": 1}, + {"ordinal": 1, "name": "b", "data_type_id": 3}, + ], + "return_data_type_id": 1, + } ) mapping[fid] = ea started = time.monotonic() failed = ImportDataTypes().execute( - FunctionDataTypesList.model_construct(items=items), + FunctionSignatures(items=items, data_types=data_types), matched_function_mapping=mapping, ) report["import_seconds"] = round(time.monotonic() - started, 3) diff --git a/tests/unit/data_types/test_import_transform.py b/tests/unit/data_types/test_import_transform.py index 0ed0362..e1b99e4 100644 --- a/tests/unit/data_types/test_import_transform.py +++ b/tests/unit/data_types/test_import_transform.py @@ -1,4 +1,3 @@ -from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -6,15 +5,9 @@ from reai_toolkit.app.transformations import import_data_types as mod from reai_toolkit.app.transformations.import_data_types import ( APPLY_CHUNK_SIZE, + FunctionSignatures, ImportDataTypes, -) -from revengai import ( - Enumeration, - FunctionDataTypesList, - FunctionInfo, - Structure, - StructureMember, - TypeDefinition, + normalise_type, ) @@ -23,35 +16,73 @@ def deci(mocker): instance = MagicMock() instance.art_lifter.lift_addr.side_effect = lambda addr: addr mocker.patch.object(mod.DecompilerInterface, "discover", return_value=instance) + mocker.patch.object(ImportDataTypes, "_install_primitives") 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 _signature(function_id: int, has_signature=True, parameters=None, return_id=None): + return { + "function_id": function_id, + "function_name": f"fn_{function_id}", + "has_signature": has_signature, + "parameters": parameters or [], + "return_data_type_id": return_id, + } + + +def _param(ordinal: int, data_type_id: int, name=None): + return {"ordinal": ordinal, "data_type_id": data_type_id, "name": name} + + +def _base(data_type_id: int, name: str, size: int = 4): + return {"data_type_id": data_type_id, "name": name, "kind": "BASE", "size": size} + +def _struct(data_type_id: int, name: str, member_type_id: int, size: int = 8): + return { + "data_type_id": data_type_id, + "name": name, + "kind": "STRUCT", + "size": size, + "definition": { + "members": [ + {"name": "field0", "offset": 0, "size": 8, "data_type_id": member_type_id} + ] + }, + } -def _functions(items): - return FunctionDataTypesList.model_construct(items=items) +def _typedef(data_type_id: int, name: str, target_id: int): + return { + "data_type_id": data_type_id, + "name": name, + "kind": "TYPEDEF", + "size": 4, + "definition": {"target_data_type_id": target_id}, + } -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 _enum(data_type_id: int, name: str): + return { + "data_type_id": data_type_id, + "name": name, + "kind": "ENUM", + "size": 4, + "definition": {"values": [{"name": "RED", "value": "0"}]}, + } + + +def _sigs(items, data_types=()): + return FunctionSignatures( + items=items, + data_types={str(entry["data_type_id"]): entry for entry in data_types}, ) def test_execute_no_items_skips_everything(deci): idt = ImportDataTypes() - assert idt.execute(_functions([SimpleNamespace(function_id=1, data_types=None)])) == set() + assert idt.execute(_sigs([_signature(1, has_signature=False)])) == set() mod.DecompilerInterface.discover.assert_not_called() @@ -59,7 +90,7 @@ 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))])) + failed = idt.execute(_sigs([_signature(1)]), matched_function_mapping={1: 0x1000}) assert failed == set() apply.assert_called_once() @@ -68,13 +99,15 @@ def test_execute_skips_discover_without_dependencies(deci, mocker): def test_execute_applies_shared_dependency_once(deci, mocker): mocker.patch.object(ImportDataTypes, "apply_function_type", return_value=True) - shared = _struct("SharedStruct") + types = [_base(1, "int"), _struct(2, "SharedStruct", 1)] items = [ - _item(1, func_types=MagicMock(addr=0x1000), func_deps=[shared]), - _item(2, func_types=MagicMock(addr=0x2000), func_deps=[shared]), + _signature(10, parameters=[_param(0, 2)]), + _signature(11, parameters=[_param(0, 2)]), ] - ImportDataTypes().execute(_functions(items)) + ImportDataTypes().execute( + _sigs(items, types), matched_function_mapping={10: 0x1000, 11: 0x2000} + ) struct_writes = [c for c in deci.structs.mock_calls if "__setitem__" in str(c)] assert len(struct_writes) == 1 @@ -82,12 +115,15 @@ def test_execute_applies_shared_dependency_once(deci, mocker): 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])] + types = [ + _base(1, "int"), + _typedef(2, "td_t", 1), + _struct(3, "Parent", 2), + _enum(4, "Colors"), + ] + items = [_signature(10, parameters=[_param(0, 3), _param(1, 4)])] - ImportDataTypes().execute(_functions(items)) + ImportDataTypes().execute(_sigs(items, types), matched_function_mapping={10: 0x1000}) writes = [c[0] for c in deci.mock_calls if "__setitem__" in c[0]] assert writes.index("typedefs.__setitem__") < writes.index("structs.__setitem__") @@ -97,23 +133,46 @@ def test_execute_applies_subdependency_before_parent(deci, mocker): 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])] + items = [_signature(10, parameters=[_param(0, 4)])] - failed = ImportDataTypes().execute(_functions(items)) + failed = ImportDataTypes().execute( + _sigs(items, [_enum(4, "Colors")]), matched_function_mapping={10: 0x1000} + ) assert failed == set() apply.assert_called_once() +def test_self_referential_struct_terminates(deci, mocker): + mocker.patch.object(ImportDataTypes, "apply_function_type", return_value=True) + pointer = { + "data_type_id": 2, + "name": "Node *", + "kind": "POINTER", + "size": 8, + "definition": {"pointee_data_type_id": 1}, + } + node = _struct(1, "Node", 2) + items = [_signature(10, parameters=[_param(0, 1)])] + + failed = ImportDataTypes().execute( + _sigs(items, [node, pointer]), matched_function_mapping={10: 0x1000} + ) + + assert failed == set() + struct_writes = [c for c in deci.structs.mock_calls if "__setitem__" in str(c)] + assert len(struct_writes) == 1 + + 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 + ImportDataTypes, "apply_function_type", side_effect=lambda proto, ea: ea % 2 == 0 ) - items = [_item(fid, func_types=MagicMock(addr=fid)) for fid in range(total)] + items = [_signature(fid) for fid in range(total)] + mapping = {fid: fid for fid in range(total)} - failed = ImportDataTypes().execute(_functions(items)) + failed = ImportDataTypes().execute(_sigs(items), matched_function_mapping=mapping) assert apply.call_count == total assert failed == {fid for fid in range(total) if fid % 2 == 1} @@ -122,31 +181,73 @@ def test_execute_chunks_and_aggregates_failures(deci, mocker): def test_execute_uses_mapping_and_fails_unmapped(deci, mocker): seen: list[int] = [] - def record(func, ea): + def record(proto, 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)), - ] + items = [_signature(1), _signature(2)] - failed = ImportDataTypes().execute(_functions(items), matched_function_mapping={1: 0x9000}) + failed = ImportDataTypes().execute(_sigs(items), matched_function_mapping={1: 0x9000}) assert seen == [0x9000] assert failed == {2} -def test_execute_skips_items_without_func_types(deci, mocker): +def test_execute_skips_items_without_signature(deci, mocker): apply = mocker.patch.object(ImportDataTypes, "apply_function_type", return_value=True) - failed = ImportDataTypes().execute(_functions([_item(1, func_types=None)])) + failed = ImportDataTypes().execute(_sigs([_signature(1, has_signature=False)])) assert failed == set() apply.assert_not_called() +def test_prototype_orders_args_by_ordinal_and_names_unnamed(deci): + resolver = mod._TypeResolver({str(1): _base(1, "int")}) + item = _signature( + 10, parameters=[_param(1, 1, name="second"), _param(0, 1)], return_id=1 + ) + + proto = ImportDataTypes._prototype(item, resolver) + + assert [a.offset for a in proto.args] == [0, 1] + assert [a.name for a in proto.args] == ["a1", "second"] + assert proto.return_type == "int" + + +def test_unresolvable_data_type_id_yields_none(): + resolver = mod._TypeResolver({}) + + assert resolver.resolve(99) == (None, None) + assert resolver.resolve(None) == (None, None) + assert resolver.artifacts == {} + + +def test_enum_value_out_of_int64_range_is_kept(): + entry = { + "data_type_id": 1, + "name": "Big", + "kind": "ENUM", + "size": 8, + "definition": { + "values": [ + {"name": "HUGE", "value": "18446744073709551615"}, + {"name": "NEG", "value": "-1"}, + {"name": "BAD", "value": "not-a-number"}, + ] + }, + } + resolver = mod._TypeResolver({"1": entry}) + + resolver.resolve(1) + + assert resolver.artifacts["Big"].members == { + "HUGE": 18446744073709551615, + "NEG": -1, + } + + _HASH = "259156281adba01eb86070f77a039e7054f268c973326adcee5fe4533f14b292" @@ -159,61 +260,19 @@ def test_execute_skips_items_without_func_types(deci, mocker): ("DWARF/stdint.h::uint32_t", "uint32_t"), ("std::vector", "std::vector"), ("int", "int"), + ("qword", "qword"), ], ) 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 *" + assert normalise_type(raw) == expected -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 - ) +def test_ghidra_primitive_declarators_are_well_formed(): + names = [ + mod._primitive_name(d) for d in mod._GHIDRA_PRIMITIVE_DECLARATORS + ] - mod.DecompilerInterface.discover.assert_not_called() - deci.functions.__setitem__.assert_not_called() + assert len(names) == len(set(names)) + assert "qword" in names + assert "undefined6" in names + assert all(name.isidentifier() for name in names) diff --git a/tests/unit/data_types/test_sdk_schemas.py b/tests/unit/data_types/test_sdk_schemas.py index 8ac9a9a..58eedef 100644 --- a/tests/unit/data_types/test_sdk_schemas.py +++ b/tests/unit/data_types/test_sdk_schemas.py @@ -1,49 +1,76 @@ -"""Schema-shape assertions for the data-types models the plugin imports. - -These guard the names/fields exercised by -reai_toolkit.app.transformations.import_data_types — the surface that broke on -the 3.93 -> 3.96 SDK bump (FunctionInfoOutput -> FunctionInfo). -""" +import inspect from revengai import ( - Argument, - Enumeration, - FunctionDataTypesList, - FunctionDataTypesListItem, - FunctionHeader, - FunctionInfo, - V2FunctionInfoFuncDepsInner, - FunctionType, - Structure, - TypeDefinition, + ArrayDataType, + BatchFunctionSignatureEntry, + DataTypeEnumValueEntry, + DataTypeMemberEntry, + DataTypesApi, + EnumDataType, + ListFunctionSignaturesOutputBody, + PointerDataType, + SignatureParameterEntry, + StructDataType, + TypedefDataType, + UnionDataType, ) -def test_function_info_has_func_types_and_deps(): - assert {"func_types", "func_deps"} <= set(FunctionInfo.model_fields) +def test_data_types_api_exposes_the_methods_the_plugin_calls(): + for method in ( + "v3_list_function_signatures", + "v3_list_analysis_data_types", + "v3_create_analysis_data_types", + "v3_update_analysis_data_types", + "v3_update_function_signature", + ): + assert callable(getattr(DataTypesApi, method)) + +def test_list_function_signatures_accepts_plugin_kwargs(): + params = inspect.signature(DataTypesApi.v3_list_function_signatures).parameters -def test_function_type_fields(): - assert {"addr", "size", "header", "name", "type"} <= set(FunctionType.model_fields) + assert {"function_ids", "include_data_types"} <= set(params) -def test_func_deps_inner_exposes_actual_instance(): - assert "actual_instance" in V2FunctionInfoFuncDepsInner.model_fields +def test_signatures_response_carries_items_and_data_types(): + assert {"items", "data_types"} <= set(ListFunctionSignaturesOutputBody.model_fields) -def test_dependency_models_expose_name_and_type(): - assert {"name", "size", "members"} <= set(Structure.model_fields) - assert {"name", "members"} <= set(Enumeration.model_fields) - assert {"name", "type"} <= set(TypeDefinition.model_fields) +def test_signature_entry_fields(): + assert { + "analysis_id", + "function_id", + "function_name", + "has_signature", + "parameters", + "return_data_type_id", + } <= set(BatchFunctionSignatureEntry.model_fields) -def test_function_header_and_argument_fields(): - assert {"name", "type", "args"} <= set(FunctionHeader.model_fields) - assert {"offset", "name", "type"} <= set(Argument.model_fields) +def test_signature_parameter_is_typed_by_id_and_ordered_by_ordinal(): + assert {"ordinal", "name", "data_type_id"} <= set( + SignatureParameterEntry.model_fields + ) -def test_data_types_list_item_fields(): - assert {"completed", "status", "data_types", "function_id"} <= set( - FunctionDataTypesListItem.model_fields +def test_record_types_expose_members_via_definition(): + for model in (StructDataType, UnionDataType): + assert {"data_type_id", "name", "kind", "size", "definition"} <= set( + model.model_fields + ) + assert {"name", "offset", "size", "data_type_id"} <= set( + DataTypeMemberEntry.model_fields ) - assert "items" in FunctionDataTypesList.model_fields + + +def test_indirect_types_reference_their_target_by_id(): + assert "definition" in TypedefDataType.model_fields + assert "definition" in PointerDataType.model_fields + assert "definition" in ArrayDataType.model_fields + + +def test_enum_values_are_strings_to_survive_wide_and_negative_values(): + assert {"name", "value"} <= set(DataTypeEnumValueEntry.model_fields) + assert "definition" in EnumDataType.model_fields + assert DataTypeEnumValueEntry.model_fields["value"].annotation is str diff --git a/tests/unit/data_types/test_service.py b/tests/unit/data_types/test_service.py index f2b03b7..9b046d4 100644 --- a/tests/unit/data_types/test_service.py +++ b/tests/unit/data_types/test_service.py @@ -2,8 +2,6 @@ import pytest from revengai.exceptions import ForbiddenException, NotFoundException -from revengai.models.function_data_types_list import FunctionDataTypesList -from revengai.models.function_data_types_list_item import FunctionDataTypesListItem from reai_toolkit.app.services.data_types import data_types_service as svc_mod from reai_toolkit.app.services.data_types.data_types_service import ( @@ -20,56 +18,52 @@ def service(): @pytest.fixture def sdk(mocker): mocker.patch.object(ImportDataTypesService, "yield_api_client") - api_class = mocker.patch.object(svc_mod, "FunctionsDataTypesApi") + api_class = mocker.patch.object(svc_mod, "DataTypesApi") api_inst = MagicMock() api_class.return_value = api_inst return api_inst -def _item(function_id: int) -> FunctionDataTypesListItem: - return FunctionDataTypesListItem.model_construct( - completed=True, - status="success", - data_types=None, - data_types_version=1, - function_id=function_id, - ) - - -def _response(items): +def _response(function_ids, data_types=(), has_signature=True): resp = MagicMock() - resp.status = True - resp.data = FunctionDataTypesList.model_construct( - total_count=len(items), - total_data_types_count=len(items), - items=items, - ) + resp.to_dict.return_value = { + "items": [ + { + "function_id": fid, + "function_name": f"fn_{fid}", + "has_signature": has_signature, + "parameters": [], + "return_data_type_id": None, + } + for fid in function_ids + ], + "data_types": [{"analysis_id": 1, "items": list(data_types)}], + } return resp def test_single_batch_when_under_threshold(service, sdk): ids = list(range(1, 11)) - sdk.list_function_data_types_for_functions.return_value = _response( - [_item(i) for i in ids] - ) + sdk.v3_list_function_signatures.return_value = _response(ids) result = service._get_data_types(ids) - assert sdk.list_function_data_types_for_functions.call_count == 1 - sdk.list_function_data_types_for_functions.assert_called_once_with(function_ids=ids) - assert [item.function_id for item in result.items] == ids + sdk.v3_list_function_signatures.assert_called_once_with( + function_ids=ids, include_data_types=True + ) + assert [item["function_id"] for item in result.items] == ids def test_chunks_large_id_list_to_avoid_uri_too_large(service, sdk): total = FUNCTION_IDS_BATCH_SIZE * 2 + 15 ids = list(range(total)) - sdk.list_function_data_types_for_functions.side_effect = lambda function_ids: _response( - [_item(i) for i in function_ids] + sdk.v3_list_function_signatures.side_effect = ( + lambda function_ids, include_data_types: _response(function_ids) ) result = service._get_data_types(ids) - calls = sdk.list_function_data_types_for_functions.call_args_list + calls = sdk.v3_list_function_signatures.call_args_list assert len(calls) == 3 assert [len(c.kwargs["function_ids"]) for c in calls] == [ FUNCTION_IDS_BATCH_SIZE, @@ -78,28 +72,45 @@ def test_chunks_large_id_list_to_avoid_uri_too_large(service, sdk): ] for c in calls: assert len(c.kwargs["function_ids"]) <= FUNCTION_IDS_BATCH_SIZE - assert [item.function_id for item in result.items] == ids + assert [item["function_id"] for item in result.items] == ids def test_empty_list_returns_none_without_calling_sdk(service, sdk): assert service._get_data_types([]) is None - sdk.list_function_data_types_for_functions.assert_not_called() + sdk.v3_list_function_signatures.assert_not_called() + +def test_data_types_are_merged_across_chunks_by_id(service, sdk): + def respond(function_ids, include_data_types): + return _response( + function_ids, + data_types=[ + {"data_type_id": fid, "name": f"t{fid}", "kind": "BASE", "size": 4} + for fid in function_ids + ], + ) -def test_skips_unsuccessful_chunks(service, sdk): - ok = _response([_item(1), _item(2)]) - bad = MagicMock(status=False, data=None) - sdk.list_function_data_types_for_functions.side_effect = [ok, bad, ok] + sdk.v3_list_function_signatures.side_effect = respond + ids = list(range(FUNCTION_IDS_BATCH_SIZE + 5)) - result = service._get_data_types( - list(range(FUNCTION_IDS_BATCH_SIZE * 2 + 1)) + result = service._get_data_types(ids) + + assert set(result.data_types) == {str(i) for i in ids} + assert result.data_types["3"]["name"] == "t3" + + +def test_data_types_without_an_id_are_skipped(service, sdk): + sdk.v3_list_function_signatures.return_value = _response( + [1], data_types=[{"name": "nameless", "kind": "BASE"}] ) - assert [item.function_id for item in result.items] == [1, 2, 1, 2] + result = service._get_data_types([1]) + + assert result.data_types == {} def test_import_data_types_computes_absent_when_remote_types_missing(service, sdk, mocker): - sdk.list_function_data_types_for_functions.return_value = _response([_item(1)]) + sdk.v3_list_function_signatures.return_value = _response([1], has_signature=False) apply = mocker.patch.object(svc_mod.ImportDataTypes, "execute", return_value=set()) result = service.import_data_types({1: 0x1000}) @@ -111,14 +122,7 @@ def test_import_data_types_computes_absent_when_remote_types_missing(service, sd def test_import_data_types_marks_apply_failures(service, sdk, mocker): - present = FunctionDataTypesListItem.model_construct( - completed=True, - status="success", - data_types=MagicMock(), - data_types_version=1, - function_id=1, - ) - sdk.list_function_data_types_for_functions.return_value = _response([present]) + sdk.v3_list_function_signatures.return_value = _response([1]) mocker.patch.object(svc_mod.ImportDataTypes, "execute", return_value={1}) result = service.import_data_types({1: 0x1000}) @@ -133,11 +137,11 @@ def test_import_data_types_empty_matches(service, sdk): assert result.error is None assert result.remote_absent_ids == set() assert result.apply_failed_ids == set() - sdk.list_function_data_types_for_functions.assert_not_called() + sdk.v3_list_function_signatures.assert_not_called() def test_import_data_types_returns_error_on_forbidden(service, sdk, mocker): - sdk.list_function_data_types_for_functions.side_effect = ForbiddenException( + sdk.v3_list_function_signatures.side_effect = ForbiddenException( status=403, reason="Forbidden" ) apply = mocker.patch.object(svc_mod.ImportDataTypes, "execute") @@ -150,11 +154,11 @@ def test_import_data_types_returns_error_on_forbidden(service, sdk, mocker): assert "Forbidden" in result.error assert result.remote_absent_ids == set() apply.assert_not_called() - assert sdk.list_function_data_types_for_functions.call_count == 1 + assert sdk.v3_list_function_signatures.call_count == 1 def test_import_data_types_treats_not_found_as_all_absent(service, sdk, mocker): - sdk.list_function_data_types_for_functions.side_effect = NotFoundException( + sdk.v3_list_function_signatures.side_effect = NotFoundException( status=404, reason="Not Found" ) apply = mocker.patch.object(svc_mod.ImportDataTypes, "execute") @@ -167,7 +171,7 @@ def test_import_data_types_treats_not_found_as_all_absent(service, sdk, mocker): def test_import_data_types_returns_error_on_unexpected_exception(service, sdk, mocker): - sdk.list_function_data_types_for_functions.side_effect = RuntimeError("boom") + sdk.v3_list_function_signatures.side_effect = RuntimeError("boom") apply = mocker.patch.object(svc_mod.ImportDataTypes, "execute") result = service.import_data_types({1: 0x1000}) From 4bffb57185ca1dd4ac9d90172a0dd55220b5ac11 Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Wed, 9 Sep 2026 12:22:26 +0100 Subject: [PATCH 03/12] feat: data-types push path --- reai_toolkit/app/app.py | 8 +- .../services/data_types/analysis_catalogue.py | 268 ++++++++++ .../variable_sync/variable_sync_service.py | 348 ++++--------- .../data_types/test_analysis_catalogue.py | 263 ++++++++++ tests/unit/variable_sync/test_service.py | 473 ++++++------------ 5 files changed, 789 insertions(+), 571 deletions(-) create mode 100644 reai_toolkit/app/services/data_types/analysis_catalogue.py create mode 100644 tests/unit/data_types/test_analysis_catalogue.py diff --git a/reai_toolkit/app/app.py b/reai_toolkit/app/app.py index 9f77cd9..1a6a714 100644 --- a/reai_toolkit/app/app.py +++ b/reai_toolkit/app/app.py @@ -15,6 +15,7 @@ from reai_toolkit.app.services.matching.matching_service import MatchingService from reai_toolkit.app.services.rename.rename_service import RenameService from reai_toolkit.app.services.upload.upload_service import UploadService +from reai_toolkit.app.services.data_types.analysis_catalogue import AnalysisDataTypesService from reai_toolkit.app.services.data_types.data_types_service import ImportDataTypesService from reai_toolkit.app.services.variable_sync.variable_sync_service import VariableSyncService @@ -52,9 +53,14 @@ def __init__(self, ida_version: str = "UNKNOWN", plugin_version: str = "UNKNOWN" self.rename_service = RenameService( netstore_service=self.netstore_service, sdk_config=sdk_config ) - self.variable_sync_service = VariableSyncService( + self.data_types_catalogue = AnalysisDataTypesService( netstore_service=self.netstore_service, sdk_config=sdk_config ) + self.variable_sync_service = VariableSyncService( + netstore_service=self.netstore_service, + sdk_config=sdk_config, + data_types_catalogue=self.data_types_catalogue, + ) self.analysis_sync_service = AnalysisSyncService( data_types_service=self.data_types_service, rename_service=self.rename_service, diff --git a/reai_toolkit/app/services/data_types/analysis_catalogue.py b/reai_toolkit/app/services/data_types/analysis_catalogue.py new file mode 100644 index 0000000..137a2ee --- /dev/null +++ b/reai_toolkit/app/services/data_types/analysis_catalogue.py @@ -0,0 +1,268 @@ +from dataclasses import dataclass, field +from typing import Any, Iterable, Optional + +from libbs.artifacts import Enum, Struct, Typedef +from loguru import logger + +from revengai import ( + Configuration, + CreateAnalysisDataTypesInputBody, + DataTypesApi, + UpdateAnalysisDataTypesInputBody, +) + +from reai_toolkit.app.core.netstore_service import SimpleNetStore +from reai_toolkit.app.interfaces.base_service import BaseService + +PAGE_SIZE = 500 +MAX_PAGES = 200 +WRITE_BATCH_SIZE = 100 + +KIND_STRUCT = "STRUCT" +KIND_ENUM = "ENUM" +KIND_TYPEDEF = "TYPEDEF" + +_DEFINED_KINDS = frozenset({KIND_STRUCT, KIND_ENUM, KIND_TYPEDEF}) + +TypeKey = tuple[str, str, str] + + +def type_key(namespace: Optional[str], name: str, kind: str) -> TypeKey: + return (namespace or "", name, (kind or "").upper()) + + +def kind_of(artifact: Any) -> Optional[str]: + if isinstance(artifact, Struct): + return KIND_STRUCT + if isinstance(artifact, Enum): + return KIND_ENUM + if isinstance(artifact, Typedef): + return KIND_TYPEDEF + return None + + +@dataclass +class Catalogue: + by_id: dict[int, dict] = field(default_factory=dict) + id_by_key: dict[TypeKey, int] = field(default_factory=dict) + id_by_name: dict[str, int] = field(default_factory=dict) + + @classmethod + def of(cls, entries: Iterable[dict]) -> "Catalogue": + catalogue = cls() + catalogue.absorb(entries) + return catalogue + + def absorb(self, entries: Iterable[dict]) -> None: + for entry in entries: + data_type_id = entry.get("data_type_id") + name = entry.get("name") + if data_type_id is None or not name: + continue + self.by_id[data_type_id] = entry + self.id_by_key[type_key(entry.get("namespace"), name, entry.get("kind"))] = data_type_id + self.id_by_name.setdefault(name, data_type_id) + + def id_of(self, key: TypeKey) -> Optional[int]: + return self.id_by_key.get(key) + + def id_of_name(self, name: str) -> Optional[int]: + return self.id_by_name.get(name) + + def __len__(self) -> int: + return len(self.by_id) + + +class AnalysisDataTypesService(BaseService): + def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration) -> None: + super().__init__(netstore_service=netstore_service, sdk_config=sdk_config) + self._cache: dict[int, Catalogue] = {} + + def invalidate(self, analysis_id: Optional[int] = None) -> None: + if analysis_id is None: + self._cache.clear() + else: + self._cache.pop(analysis_id, None) + + def catalogue(self, analysis_id: int) -> Catalogue: + cached = self._cache.get(analysis_id) + if cached is None: + cached = self.sync(analysis_id) + return cached + + def sync(self, analysis_id: int) -> Catalogue: + entries: list[dict] = [] + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + client = DataTypesApi(api_client) + offset = 0 + for _ in range(MAX_PAGES): + page = client.v3_list_analysis_data_types( + analysis_id=analysis_id, offset=offset, limit=PAGE_SIZE + ).to_dict() + items = page.get("items") or [] + entries.extend(items) + if len(items) < PAGE_SIZE: + break + offset += len(items) + else: + logger.warning( + f"RevEng.AI: stopped paging analysis {analysis_id} data types at {MAX_PAGES} pages" + ) + + catalogue = Catalogue.of(entries) + self._cache[analysis_id] = catalogue + logger.debug(f"RevEng.AI: analysis {analysis_id} holds {len(catalogue)} data types") + return catalogue + + def resolve(self, analysis_id: int, type_name: Optional[str]) -> Optional[int]: + if not type_name: + return None + return self.catalogue(analysis_id).id_of_name(type_name) + + def ensure(self, analysis_id: int, artifacts: dict[str, Any]) -> dict[str, int]: + if not artifacts: + return {} + + catalogue = self.catalogue(analysis_id) + ids: dict[str, int] = {} + missing: dict[str, Any] = {} + + for name, artifact in artifacts.items(): + kind = kind_of(artifact) + if kind is None: + continue + existing = catalogue.id_of(type_key("", name, kind)) + if existing is None: + missing[name] = artifact + else: + ids[name] = existing + + if missing: + ids.update(self._create_placeholders(analysis_id, missing, catalogue)) + + defined = { + name: artifacts[name] + for name in ids + if kind_of(artifacts[name]) in _DEFINED_KINDS + } + if defined: + self._write_definitions(analysis_id, defined, ids, catalogue) + + return ids + + def _create_placeholders( + self, analysis_id: int, missing: dict[str, Any], catalogue: Catalogue + ) -> dict[str, int]: + bodies = [ + { + "kind": kind_of(artifact), + "name": name, + "namespace": "", + "size": getattr(artifact, "size", None) or 0, + } + for name, artifact in missing.items() + ] + + created: dict[str, int] = {} + for chunk in _chunks(bodies, WRITE_BATCH_SIZE): + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + client = DataTypesApi(api_client) + out = client.v3_create_analysis_data_types( + analysis_id=analysis_id, + create_analysis_data_types_input_body=CreateAnalysisDataTypesInputBody.model_construct( + data_types=chunk + ), + ).to_dict() + + entries = out.get("data_types") or [] + catalogue.absorb(entries) + for entry in entries: + name = entry.get("name") + data_type_id = entry.get("data_type_id") + if name is not None and data_type_id is not None: + created[name] = data_type_id + + return created + + def _write_definitions( + self, + analysis_id: int, + artifacts: dict[str, Any], + ids: dict[str, int], + catalogue: Catalogue, + ) -> None: + bodies = [] + for name, artifact in artifacts.items(): + definition = self._definition(artifact, analysis_id, ids, catalogue) + if definition is None: + continue + bodies.append( + { + "data_type_id": ids[name], + "kind": kind_of(artifact), + "name": name, + "namespace": "", + "size": getattr(artifact, "size", None) or 0, + "definition": definition, + } + ) + + for chunk in _chunks(bodies, WRITE_BATCH_SIZE): + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + client = DataTypesApi(api_client) + out = client.v3_update_analysis_data_types( + analysis_id=analysis_id, + update_analysis_data_types_input_body=UpdateAnalysisDataTypesInputBody.model_construct( + data_types=chunk + ), + ).to_dict() + catalogue.absorb(out.get("data_types") or []) + + def _definition( + self, + artifact: Any, + analysis_id: int, + ids: dict[str, int], + catalogue: Catalogue, + ) -> Optional[dict]: + if isinstance(artifact, Struct): + members = [] + for member in (artifact.members or {}).values(): + members.append( + { + "name": member.name or "", + "offset": member.offset, + "size": member.size or 0, + "is_bitfield": False, + "data_type_id": self._reference(member.type, ids, catalogue), + } + ) + return {"members": members} + + if isinstance(artifact, Enum): + return { + "values": [ + {"name": str(key), "value": str(value)} + for key, value in (artifact.members or {}).items() + ] + } + + if isinstance(artifact, Typedef): + return {"target_data_type_id": self._reference(artifact.type, ids, catalogue)} + + return None + + @staticmethod + def _reference( + type_name: Optional[str], ids: dict[str, int], catalogue: Catalogue + ) -> Optional[int]: + if not type_name: + return None + if type_name in ids: + return ids[type_name] + return catalogue.id_of_name(type_name) + + +def _chunks(items: list, size: int): + for start in range(0, len(items), size): + yield items[start:start + size] 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 93aa089..ddb9f45 100644 --- a/reai_toolkit/app/services/variable_sync/variable_sync_service.py +++ b/reai_toolkit/app/services/variable_sync/variable_sync_service.py @@ -1,31 +1,26 @@ import queue +import re import threading import time -from typing import Optional, Tuple +from typing import Any, Optional, Tuple -from libbs.artifacts import Enum, FunctionHeader, StackVariable, Struct, Typedef +from libbs.artifacts import Enum, StackVariable, Struct, Typedef 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 ( - FunctionArgument, - BaseResponseFunctionDataTypesList, - BatchUpdateDataTypesInputBody, - BatchUpdateDataTypesItem, - BatchUpdateDataTypesOutputBody, - Configuration, - FunctionDependency, - FunctionInfo, - FunctionsDataTypesApi, - FunctionType, -) -from revengai import ApiException -from revengai.models.function_header import FunctionHeader as SdkFunctionHeader -from revengai.models.function_stack_variable import FunctionStackVariable as SdkStackVariable +from revengai import ApiException, Configuration, DataTypesApi +from revengai.exceptions import NotFoundException from reai_toolkit.app.core.netstore_service import SimpleNetStore from reai_toolkit.app.interfaces.thread_service import IThreadService +from reai_toolkit.app.services.data_types.analysis_catalogue import ( + AnalysisDataTypesService, +) + +MAX_TYPE_DEPENDENCIES = 200 + +_ARRAY_SUFFIX_RE = re.compile(r"\[[^\]]*\]") @execute_read @@ -62,12 +57,18 @@ class VariableSyncService(IThreadService): _q: queue.Queue[Tuple[int, int, object]] = queue.Queue() _last_ts: dict[tuple, float] = {} _debounce_ms: int = 400 - _max_retries: int = 3 - _type_batch_size: int = 50 - def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration): + def __init__( + self, + netstore_service: SimpleNetStore, + sdk_config: Configuration, + data_types_catalogue: Optional[AnalysisDataTypesService] = None, + ): super().__init__(netstore_service=netstore_service, sdk_config=sdk_config) self._deci = None + self._catalogue = data_types_catalogue or AnalysisDataTypesService( + netstore_service=netstore_service, sdk_config=sdk_config + ) def attach_decompiler(self, deci) -> None: self._deci = deci @@ -103,7 +104,7 @@ def _worker(self, stop_event: Optional[threading.Event] = None) -> None: except queue.Empty: continue try: - self._push_change(function_id, func_addr, artifact) + self._push_change(function_id, func_addr) except ApiException as e: logger.error(f"RevEng.AI: failed to push data types: HTTP {e.status} {e.reason}") except Exception as e: @@ -111,58 +112,24 @@ def _worker(self, stop_event: Optional[threading.Event] = None) -> None: finally: self._q.task_done() - def _push_change(self, function_id: int, func_addr: int, artifact: object) -> None: + def _push_change(self, function_id: int, func_addr: int) -> None: logger.debug(f"RevEng.AI: processing data types change for function {function_id}") analysis_id: int | None = self.netstore_service.get_analysis_id() if analysis_id is None: logger.debug("RevEng.AI: no analysis id; skipping data types push") return - for _ in range(self._max_retries): - fetched = self._fetch_function_info(function_id) - if fetched is None: - # First write for this function: build the object from current state. - info = self._build_function_info(func_addr) - if info is None: - logger.debug(f"RevEng.AI: could not build data types for function {function_id}; skipping push") - return - version = 0 - else: - info, version = fetched - if not self._patch(info, artifact): - logger.debug(f"RevEng.AI: no data types change for function {function_id}; skipping push") - return - - status: str = self._push(function_id, analysis_id, info, version) - # A concurrent edit moved the stored version on; re-fetch and reapply. - if status == "version_conflict": - time.sleep(0.2) - continue - if status == "updated": - logger.info(f"RevEng.AI: pushed data types for function {function_id}") - else: - logger.warning( - f"RevEng.AI: data types update for function {function_id} returned {status}" - ) - return - - def _fetch_function_info(self, function_id: int) -> Optional[Tuple[FunctionInfo, int]]: - with self.yield_api_client(sdk_config=self.sdk_config) as api_client: - client = FunctionsDataTypesApi(api_client) - resp: BaseResponseFunctionDataTypesList = ( - client.list_function_data_types_for_functions(function_ids=[function_id]) # type: ignore + signature = self._build_signature(analysis_id, func_addr) + if signature is None: + logger.debug( + f"RevEng.AI: could not build a signature for function {function_id}; skipping push" ) + return - if not resp.status or not resp.data: - return None - - for item in resp.data.items: - if item.function_id == function_id and item.data_types is not None: - return item.data_types, (item.data_types_version or 0) - - return None + if self._put_signature(analysis_id, function_id, signature): + logger.info(f"RevEng.AI: pushed data types for function {function_id}") - def _build_function_info(self, func_addr: int) -> Optional[FunctionInfo]: + def _build_signature(self, analysis_id: int, func_addr: int) -> Optional[dict]: if self._deci is None: return None @@ -171,159 +138,88 @@ def _build_function_info(self, func_addr: int) -> Optional[FunctionInfo]: return None header = func.header - args = { - hex(offset): FunctionArgument( - offset=offset, name=arg.name or "", type=arg.type or "", size=arg.size or 0 - ) - for offset, arg in (header.args or {}).items() - } - stack_vars = { - hex(offset): SdkStackVariable( - offset=offset, - name=svar.name or "", - type=svar.type or "", - size=svar.size or 0, - addr=svar.addr or func.addr, + artifacts = self._collect_type_artifacts(func) + ids = self._catalogue.ensure(analysis_id, artifacts) if artifacts else {} + + parameters = [] + for ordinal, (_, arg) in enumerate(sorted((header.args or {}).items())): + parameters.append( + { + "ordinal": ordinal, + "name": arg.name or None, + "data_type_id": self._type_id(analysis_id, arg.type, ids), + } ) - for offset, svar in (func.stack_vars or {}).items() + + return { + "parameters": parameters, + "return_data_type_id": self._type_id(analysis_id, header.type, ids), } - func_type = FunctionType( - addr=func.addr, - size=func.size or 0, - header=SdkFunctionHeader( - name=header.name or "", addr=header.addr or func.addr, type=header.type or "", args=args - ), - stack_vars=stack_vars, - name=getattr(func, "name", None) or header.name or "", - type=header.type or "", - artifact_type="Function", - ) - return FunctionInfo(func_types=func_type, func_deps=self._collect_func_deps(func_type)) - def _collect_func_deps(self, func_type: FunctionType) -> list: - pending: list[str] = [func_type.type] - pending.extend(arg.type for arg in (func_type.header.args or {}).values()) - pending.extend(svar.type for svar in (func_type.stack_vars or {}).values()) + def _type_id( + self, analysis_id: int, type_name: Optional[str], ids: dict[str, int] + ) -> Optional[int]: + if not type_name: + return None + if type_name in ids: + return ids[type_name] + return self._catalogue.resolve(analysis_id, type_name) + + def _collect_type_artifacts(self, func: Any) -> dict[str, Any]: + header = func.header + pending: list[Optional[str]] = [header.type] + pending.extend(arg.type for arg in (header.args or {}).values()) + pending.extend(svar.type for svar in (func.stack_vars or {}).values()) - deps: dict[str, FunctionDependency] = {} + artifacts: dict[str, Any] = {} seen: set[str] = set() - while pending and len(deps) < 200: + while pending and len(artifacts) < MAX_TYPE_DEPENDENCIES: name = self._base_type_name(pending.pop()) if not name or name in seen: continue seen.add(name) - dep, referenced = self._resolve_type(name) - if dep is None: + + artifact = _read_named_type(self._deci, name) + if artifact is None: continue - deps[name] = dep - pending.extend(referenced) + artifacts[name] = artifact + pending.extend(self._referenced_types(artifact)) - return list(deps.values()) + return artifacts - def _resolve_type(self, name: str) -> Tuple[Optional[FunctionDependency], list]: - artifact = _read_named_type(self._deci, name) + @staticmethod + def _referenced_types(artifact: Any) -> list: if isinstance(artifact, Typedef): - return ( - FunctionDependency(name=artifact.name, type=artifact.type or "", artifact_type="Typedef"), - [artifact.type], - ) + return [artifact.type] if isinstance(artifact, Struct): - members = { - hex(member.offset): { - "name": member.name or "", - "offset": member.offset, - "type": member.type or "", - "size": member.size or 0, - } - for member in artifact.members.values() - } - referenced = [member.type for member in artifact.members.values()] - return ( - FunctionDependency(name=artifact.name, size=artifact.size, members=members, artifact_type="Struct"), - referenced, - ) + return [member.type for member in (artifact.members or {}).values()] if isinstance(artifact, Enum): - members = {str(key): int(value) for key, value in (artifact.members or {}).items()} - return ( - FunctionDependency(name=artifact.name, members=members, artifact_type="Enum"), - [], - ) - return None, [] + return [] + return [] @staticmethod def _base_type_name(type_str: Optional[str]) -> Optional[str]: if not type_str: return None - cleaned = type_str.replace("*", " ").replace("[", " ").replace("]", " ") + cleaned = _ARRAY_SUFFIX_RE.sub(" ", type_str).replace("*", " ") keywords = {"const", "volatile", "struct", "union", "enum", "unsigned", "signed"} tokens = [tok for tok in cleaned.split() if tok not in keywords] return tokens[-1] if tokens else None - def _push(self, function_id: int, analysis_id: int, info: FunctionInfo, version: int) -> str: - body = BatchUpdateDataTypesInputBody( - functions=[ - BatchUpdateDataTypesItem( + def _put_signature(self, analysis_id: int, function_id: int, signature: dict) -> bool: + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + DataTypesApi(api_client).v3_update_function_signature( + analysis_id=analysis_id, function_id=function_id, - data_types=info.to_dict(), - data_types_version=version, + update_function_signature_input_body=signature, ) - ] - ) - with self.yield_api_client(sdk_config=self.sdk_config) as api_client: - client = FunctionsDataTypesApi(api_client) - out: BatchUpdateDataTypesOutputBody = client.batch_update_function_data_types( - analysis_id=analysis_id, - batch_update_data_types_input_body=body, + except NotFoundException: + logger.debug( + f"RevEng.AI: no extracted signature for function {function_id}; nothing to edit" ) - - if out.results: - return out.results[0].status - return "error" - - def _patch(self, info: FunctionInfo, artifact: object) -> bool: - ft = info.func_types - if ft is None: return False - - before = ft.to_dict() - if isinstance(artifact, StackVariable): - self._patch_stack_var(ft, artifact) - elif isinstance(artifact, FunctionHeader): - self._patch_header(ft, artifact) - # Skip the push when nothing actually changed: pure function renames also - # raise a header event but are owned by the rename service, and sync-applied - # edits already match the stored blob. - return ft.to_dict() != before - - @staticmethod - def _patch_stack_var(ft: object, svar: StackVariable) -> None: - if not ft.stack_vars: - return - for entry in ft.stack_vars.values(): - if entry.offset == svar.offset: - if svar.name is not None: - entry.name = svar.name - if svar.type is not None: - entry.type = svar.type - return - - @staticmethod - def _patch_header(ft: object, fheader: FunctionHeader) -> None: - header = ft.header - if header is None: - return - if fheader.type: - header.type = fheader.type - ft.type = fheader.type - if header.args and fheader.args: - for offset, arg in fheader.args.items(): - for entry in header.args.values(): - if entry.offset == offset: - if arg.name is not None: - entry.name = arg.name - if arg.type is not None: - entry.type = arg.type - break + return True def push_local_function_types_batch(self, targets: dict[int, int], analysis_id: int | None) -> int: if not targets or analysis_id is None: @@ -339,74 +235,22 @@ def push_local_function_types_batch(self, targets: dict[int, int], analysis_id: return 0 base: int = self._deci.binary_base_addr - infos: dict[int, FunctionInfo] = {} + updated: int = 0 for function_id, ea in targets.items(): try: - info = self._build_function_info(ea - base) + signature = self._build_signature(analysis_id, ea - base) except Exception as e: logger.debug(f"RevEng.AI: could not build local types for function {function_id}: {e}") continue - if info is not None: - infos[function_id] = info - - if not infos: - return 0 - - versions: dict[int, int] = self._fetch_type_versions(list(infos.keys())) - pending: dict[int, FunctionInfo] = dict(infos) - updated: int = 0 + if signature is None: + continue - for _ in range(self._max_retries): - if not pending: - break - conflicts: dict[int, FunctionInfo] = {} - items = list(pending.items()) - for start in range(0, len(items), self._type_batch_size): - chunk = items[start:start + self._type_batch_size] - for result in self._push_types_batch(analysis_id, chunk, versions): - if result.status == "updated": - updated += 1 - elif result.status == "version_conflict": - conflicts[result.function_id] = infos[result.function_id] - else: - logger.warning( - f"RevEng.AI: type push for function {result.function_id} returned {result.status}" - ) - pending = conflicts - if pending: - versions.update(self._fetch_type_versions(list(pending.keys()))) + try: + if self._put_signature(analysis_id, function_id, signature): + updated += 1 + except ApiException as e: + logger.warning( + f"RevEng.AI: type push for function {function_id} failed: HTTP {e.status} {e.reason}" + ) return updated - - def _fetch_type_versions(self, function_ids: list[int]) -> dict[int, int]: - versions: dict[int, int] = {} - with self.yield_api_client(sdk_config=self.sdk_config) as api_client: - client = FunctionsDataTypesApi(api_client) - for start in range(0, len(function_ids), self._type_batch_size): - chunk = function_ids[start:start + self._type_batch_size] - resp: BaseResponseFunctionDataTypesList = ( - client.list_function_data_types_for_functions(function_ids=chunk) # type: ignore - ) - if resp.status and resp.data: - for item in resp.data.items: - versions[item.function_id] = item.data_types_version or 0 - return versions - - def _push_types_batch(self, analysis_id: int, items: list, versions: dict[int, int]) -> list: - body = BatchUpdateDataTypesInputBody( - functions=[ - BatchUpdateDataTypesItem( - function_id=function_id, - data_types=info.to_dict(), - data_types_version=versions.get(function_id, 0), - ) - for function_id, info in items - ] - ) - with self.yield_api_client(sdk_config=self.sdk_config) as api_client: - client = FunctionsDataTypesApi(api_client) - out: BatchUpdateDataTypesOutputBody = client.batch_update_function_data_types( - analysis_id=analysis_id, - batch_update_data_types_input_body=body, - ) - return out.results or [] diff --git a/tests/unit/data_types/test_analysis_catalogue.py b/tests/unit/data_types/test_analysis_catalogue.py new file mode 100644 index 0000000..8159d36 --- /dev/null +++ b/tests/unit/data_types/test_analysis_catalogue.py @@ -0,0 +1,263 @@ +from unittest.mock import MagicMock + +import pytest +from libbs.artifacts import Enum, Struct, StructMember, Typedef + +from reai_toolkit.app.services.data_types import analysis_catalogue as mod +from reai_toolkit.app.services.data_types.analysis_catalogue import ( + MAX_PAGES, + PAGE_SIZE, + WRITE_BATCH_SIZE, + AnalysisDataTypesService, + Catalogue, +) + +ANALYSIS = 7 + + +@pytest.fixture +def service(): + return AnalysisDataTypesService(netstore_service=MagicMock(), sdk_config=MagicMock()) + + +@pytest.fixture +def sdk(mocker): + mocker.patch.object(AnalysisDataTypesService, "yield_api_client") + api_class = mocker.patch.object(mod, "DataTypesApi") + api = MagicMock() + api_class.return_value = api + return api + + +def _entry(data_type_id, name, kind="STRUCT", namespace="", size=8): + return { + "data_type_id": data_type_id, + "name": name, + "kind": kind, + "namespace": namespace, + "size": size, + } + + +def _page(items): + resp = MagicMock() + resp.to_dict.return_value = {"items": items, "total_count": len(items)} + return resp + + +def _written(items): + resp = MagicMock() + resp.to_dict.return_value = {"data_types": items} + return resp + + +def test_catalogue_indexes_by_key_and_name(): + catalogue = Catalogue.of([_entry(1, "Foo"), _entry(2, "Bar", kind="ENUM")]) + + assert catalogue.id_of(("", "Foo", "STRUCT")) == 1 + assert catalogue.id_of(("", "Bar", "ENUM")) == 2 + assert catalogue.id_of(("", "Foo", "ENUM")) is None + assert catalogue.id_of_name("Foo") == 1 + + +def test_catalogue_keeps_data_type_id_zero(): + catalogue = Catalogue.of([_entry(0, "Zero")]) + + assert catalogue.id_of_name("Zero") == 0 + assert len(catalogue) == 1 + + +def test_sync_pages_until_short_page(service, sdk): + first = [_entry(i, f"T{i}") for i in range(PAGE_SIZE)] + second = [_entry(PAGE_SIZE, "Last")] + sdk.v3_list_analysis_data_types.side_effect = [_page(first), _page(second)] + + catalogue = service.sync(ANALYSIS) + + assert len(catalogue) == PAGE_SIZE + 1 + offsets = [c.kwargs["offset"] for c in sdk.v3_list_analysis_data_types.call_args_list] + assert offsets == [0, PAGE_SIZE] + + +def test_sync_stops_at_max_pages_when_server_never_shortens(service, sdk): + sdk.v3_list_analysis_data_types.side_effect = lambda **kw: _page( + [_entry(kw["offset"] + i, f"T{kw['offset'] + i}") for i in range(PAGE_SIZE)] + ) + + service.sync(ANALYSIS) + + assert sdk.v3_list_analysis_data_types.call_count == MAX_PAGES + + +def test_catalogue_is_cached_until_invalidated(service, sdk): + sdk.v3_list_analysis_data_types.return_value = _page([_entry(1, "Foo")]) + + service.catalogue(ANALYSIS) + service.catalogue(ANALYSIS) + assert sdk.v3_list_analysis_data_types.call_count == 1 + + service.invalidate(ANALYSIS) + service.catalogue(ANALYSIS) + assert sdk.v3_list_analysis_data_types.call_count == 2 + + +def test_ensure_resolves_before_creating(service, sdk): + sdk.v3_list_analysis_data_types.return_value = _page([_entry(11, "Known")]) + known = Struct(name="Known", size=8, members={}) + + ids = service.ensure(ANALYSIS, {"Known": known}) + + assert ids == {"Known": 11} + sdk.v3_create_analysis_data_types.assert_not_called() + + +def test_ensure_creates_only_the_gaps(service, sdk): + sdk.v3_list_analysis_data_types.return_value = _page([_entry(11, "Known")]) + sdk.v3_create_analysis_data_types.return_value = _written([_entry(12, "New")]) + sdk.v3_update_analysis_data_types.return_value = _written([]) + + ids = service.ensure( + ANALYSIS, + { + "Known": Struct(name="Known", size=8, members={}), + "New": Struct(name="New", size=4, members={}), + }, + ) + + assert ids == {"Known": 11, "New": 12} + body = sdk.v3_create_analysis_data_types.call_args.kwargs[ + "create_analysis_data_types_input_body" + ] + assert [d["name"] for d in body.data_types] == ["New"] + + +def test_create_bodies_carry_no_data_type_id(service, sdk): + sdk.v3_list_analysis_data_types.return_value = _page([]) + sdk.v3_create_analysis_data_types.return_value = _written([_entry(1, "New")]) + sdk.v3_update_analysis_data_types.return_value = _written([]) + + service.ensure(ANALYSIS, {"New": Struct(name="New", size=4, members={})}) + + body = sdk.v3_create_analysis_data_types.call_args.kwargs[ + "create_analysis_data_types_input_body" + ] + assert all("data_type_id" not in d for d in body.data_types) + + +def test_definitions_are_written_after_ids_exist(service, sdk): + sdk.v3_list_analysis_data_types.return_value = _page([_entry(1, "int", kind="BASE", size=4)]) + sdk.v3_create_analysis_data_types.return_value = _written( + [_entry(2, "Node"), _entry(3, "NodeRef", kind="TYPEDEF")] + ) + sdk.v3_update_analysis_data_types.return_value = _written([]) + + node = Struct( + name="Node", + size=8, + members={0: StructMember(name="value", offset=0, type_="int", size=4)}, + ) + alias = Typedef(name="NodeRef", type_="Node") + + service.ensure(ANALYSIS, {"Node": node, "NodeRef": alias}) + + body = sdk.v3_update_analysis_data_types.call_args.kwargs[ + "update_analysis_data_types_input_body" + ] + written = {d["name"]: d for d in body.data_types} + assert written["Node"]["definition"]["members"][0]["data_type_id"] == 1 + assert written["NodeRef"]["definition"]["target_data_type_id"] == 2 + assert all("data_type_id" in d for d in body.data_types) + + +def test_self_referential_struct_resolves_to_its_own_new_id(service, sdk): + sdk.v3_list_analysis_data_types.return_value = _page([]) + sdk.v3_create_analysis_data_types.return_value = _written([_entry(5, "Node")]) + sdk.v3_update_analysis_data_types.return_value = _written([]) + + node = Struct( + name="Node", + size=8, + members={0: StructMember(name="next", offset=0, type_="Node", size=8)}, + ) + + service.ensure(ANALYSIS, {"Node": node}) + + body = sdk.v3_update_analysis_data_types.call_args.kwargs[ + "update_analysis_data_types_input_body" + ] + assert body.data_types[0]["definition"]["members"][0]["data_type_id"] == 5 + + +def test_unresolvable_member_type_is_left_unresolved(service, sdk): + sdk.v3_list_analysis_data_types.return_value = _page([]) + sdk.v3_create_analysis_data_types.return_value = _written([_entry(1, "Foo")]) + sdk.v3_update_analysis_data_types.return_value = _written([]) + + foo = Struct( + name="Foo", + size=8, + members={0: StructMember(name="x", offset=0, type_="NoSuchType", size=8)}, + ) + + service.ensure(ANALYSIS, {"Foo": foo}) + + body = sdk.v3_update_analysis_data_types.call_args.kwargs[ + "update_analysis_data_types_input_body" + ] + assert body.data_types[0]["definition"]["members"][0]["data_type_id"] is None + + +def test_enum_values_are_sent_as_strings(service, sdk): + sdk.v3_list_analysis_data_types.return_value = _page([]) + sdk.v3_create_analysis_data_types.return_value = _written( + [_entry(1, "Colors", kind="ENUM", size=4)] + ) + sdk.v3_update_analysis_data_types.return_value = _written([]) + + service.ensure(ANALYSIS, {"Colors": Enum(name="Colors", members={"RED": 0, "BIG": 2**64 - 1})}) + + body = sdk.v3_update_analysis_data_types.call_args.kwargs[ + "update_analysis_data_types_input_body" + ] + values = body.data_types[0]["definition"]["values"] + assert {v["name"]: v["value"] for v in values} == { + "RED": "0", + "BIG": "18446744073709551615", + } + + +def test_writes_are_chunked(service, sdk): + total = WRITE_BATCH_SIZE + 10 + sdk.v3_list_analysis_data_types.return_value = _page([]) + sdk.v3_create_analysis_data_types.side_effect = lambda **kw: _written( + [ + _entry(i, d["name"]) + for i, d in enumerate( + kw["create_analysis_data_types_input_body"].data_types, start=1 + ) + ] + ) + sdk.v3_update_analysis_data_types.return_value = _written([]) + + artifacts = { + f"S{i}": Struct(name=f"S{i}", size=4, members={}) for i in range(total) + } + service.ensure(ANALYSIS, artifacts) + + sizes = [ + len(c.kwargs["create_analysis_data_types_input_body"].data_types) + for c in sdk.v3_create_analysis_data_types.call_args_list + ] + assert sizes == [WRITE_BATCH_SIZE, 10] + + +def test_ensure_ignores_artifacts_of_unknown_kind(service, sdk): + sdk.v3_list_analysis_data_types.return_value = _page([]) + + assert service.ensure(ANALYSIS, {"weird": object()}) == {} + sdk.v3_create_analysis_data_types.assert_not_called() + + +def test_ensure_with_no_artifacts_touches_nothing(service, sdk): + assert service.ensure(ANALYSIS, {}) == {} + sdk.v3_list_analysis_data_types.assert_not_called() diff --git a/tests/unit/variable_sync/test_service.py b/tests/unit/variable_sync/test_service.py index 4621024..e561b25 100644 --- a/tests/unit/variable_sync/test_service.py +++ b/tests/unit/variable_sync/test_service.py @@ -1,31 +1,48 @@ import queue -from types import SimpleNamespace from unittest.mock import MagicMock import pytest -from libbs.artifacts import Function, FunctionArgument, FunctionHeader, StackVariable -from revengai import ( - Argument, - FunctionInfo, - FunctionType, +from libbs.artifacts import ( + Function, + FunctionArgument, + FunctionHeader, + StackVariable, + Struct, + StructMember, + Typedef, ) -from revengai.models.function_header import FunctionHeader as SdkFunctionHeader -from revengai.models.stack_variable import StackVariable as SdkStackVariable +from revengai.exceptions import NotFoundException from reai_toolkit.app.services.variable_sync import variable_sync_service as svc_mod from reai_toolkit.app.services.variable_sync.variable_sync_service import ( VariableSyncService, ) +ANALYSIS = 7 + @pytest.fixture def netstore(): - return MagicMock() + store = MagicMock() + store.get_analysis_id.return_value = ANALYSIS + return store @pytest.fixture -def service(netstore): - svc = VariableSyncService(netstore_service=netstore, sdk_config=MagicMock()) +def catalogue(): + cat = MagicMock() + cat.ensure.return_value = {} + cat.resolve.return_value = None + return cat + + +@pytest.fixture +def service(netstore, catalogue): + svc = VariableSyncService( + netstore_service=netstore, + sdk_config=MagicMock(), + data_types_catalogue=catalogue, + ) VariableSyncService._q = queue.Queue() VariableSyncService._last_ts = {} return svc @@ -34,376 +51,196 @@ def service(netstore): @pytest.fixture def sdk(mocker): mocker.patch.object(VariableSyncService, "yield_api_client") - api_class = mocker.patch.object(svc_mod, "FunctionsDataTypesApi") - api_inst = MagicMock() - api_class.return_value = api_inst - return api_inst + api_class = mocker.patch.object(svc_mod, "DataTypesApi") + api = MagicMock() + api_class.return_value = api + return api -def _sdk_stack_var(offset: int, name: str, type_: str) -> SdkStackVariable: - return SdkStackVariable.model_construct( - last_change=None, offset=offset, name=name, type=type_, size=8, addr=0x1000 - ) - - -def _sdk_arg(offset: int, name: str, type_: str) -> Argument: - return Argument.model_construct( - last_change=None, offset=offset, name=name, type=type_, size=8 - ) - - -def _function_info(stack_vars=None, args=None, ret_type="void") -> FunctionInfo: - header = SdkFunctionHeader.model_construct( - last_change=None, - name="f", - addr=0x1000, - type=ret_type, - args=args or {}, - ) - func_types = FunctionType.model_construct( - last_change=None, +def _function(args=None, ret="int", stack_vars=None): + return Function( addr=0x1000, - size=10, - header=header, + header=FunctionHeader(name="f", addr=0x1000, type_=ret, args=args or {}), stack_vars=stack_vars or {}, - name="f", - type=ret_type, - artifact_type="Function", ) - return FunctionInfo.model_construct(func_types=func_types, func_deps=[]) - - -def test_patch_stack_var_matches_by_offset(service): - info = _function_info(stack_vars={"-0x20": _sdk_stack_var(-32, "local_20", "char")}) - changed = service._patch(info, StackVariable(stack_offset=-32, name="counter", type_="int", size=4, addr=0x0)) - assert changed is True - entry = info.func_types.stack_vars["-0x20"] - assert entry.name == "counter" - assert entry.type == "int" +def _arg(offset, name, type_): + return FunctionArgument(offset=offset, name=name, type_=type_, size=8) -def test_patch_stack_var_no_offset_match_is_noop(service): - info = _function_info(stack_vars={"-0x20": _sdk_stack_var(-32, "local_20", "char")}) - - changed = service._patch(info, StackVariable(stack_offset=-99, name="x", type_="int", size=4, addr=0x0)) +def _attach(service, mocker, func): + service.attach_decompiler(MagicMock()) + mocker.patch.object(svc_mod, "_read_decompiler_function", return_value=func) - assert changed is False - assert info.func_types.stack_vars["-0x20"].name == "local_20" +def test_push_change_no_analysis_id_does_nothing(service, sdk, netstore): + netstore.get_analysis_id.return_value = None -def test_patch_header_updates_arg_and_return_type(service): - info = _function_info(args={"0x0": _sdk_arg(0, "a1", "int")}, ret_type="void") - fheader = FunctionHeader( - name="f", addr=0x0, type_="int", args={0: FunctionArgument(offset=0, name="count", type_="size_t", size=8)} - ) + service._push_change(1, 0x1000) - changed = service._patch(info, fheader) + sdk.v3_update_function_signature.assert_not_called() - assert changed is True - assert info.func_types.header.args["0x0"].name == "count" - assert info.func_types.header.args["0x0"].type == "size_t" - assert info.func_types.header.type == "int" - assert info.func_types.type == "int" +def test_push_change_skips_when_no_decompiler(service, sdk): + service._push_change(1, 0x1000) -def test_push_local_types_batch_noop_without_targets_or_analysis(service): - service._deci = MagicMock() - assert service.push_local_function_types_batch({}, analysis_id=1) == 0 - assert service.push_local_function_types_batch({1: 0x1000}, analysis_id=None) == 0 + sdk.v3_update_function_signature.assert_not_called() -def test_push_local_types_batch_uses_version_zero_when_absent(service, sdk, mocker): - service._deci = MagicMock() - service._deci.binary_base_addr = 0x400000 - build = mocker.patch.object(service, "_build_function_info", return_value=_function_info()) - sdk.list_function_data_types_for_functions.return_value = MagicMock( - status=True, data=MagicMock(items=[]) - ) - out = MagicMock() - out.results = [SimpleNamespace(function_id=1, status="updated", data_types_version=1)] - sdk.batch_update_function_data_types.return_value = out +def test_push_change_puts_signature(service, sdk, mocker, catalogue): + catalogue.resolve.side_effect = lambda _a, name: {"int": 1, "char *": 2}.get(name) + _attach(service, mocker, _function(args={0: _arg(0, "count", "int")}, ret="int")) - updated = service.push_local_function_types_batch({1: 0x401000}, analysis_id=9) + service._push_change(42, 0x1000) - assert updated == 1 - build.assert_called_once_with(0x401000 - 0x400000) - body = sdk.batch_update_function_data_types.call_args.kwargs[ - "batch_update_data_types_input_body" - ] - assert body.functions[0].function_id == 1 - assert body.functions[0].data_types_version == 0 + call = sdk.v3_update_function_signature.call_args.kwargs + assert call["analysis_id"] == ANALYSIS + assert call["function_id"] == 42 + body = call["update_function_signature_input_body"] + assert body["return_data_type_id"] == 1 + assert body["parameters"] == [{"ordinal": 0, "name": "count", "data_type_id": 1}] -def test_push_local_types_batch_overwrites_with_remote_version(service, sdk, mocker): - service._deci = MagicMock() - service._deci.binary_base_addr = 0 - mocker.patch.object(service, "_build_function_info", return_value=_function_info()) - sdk.list_function_data_types_for_functions.return_value = MagicMock( - status=True, - data=MagicMock(items=[SimpleNamespace(function_id=1, data_types_version=5)]), - ) - out = MagicMock() - out.results = [SimpleNamespace(function_id=1, status="updated", data_types_version=6)] - sdk.batch_update_function_data_types.return_value = out +def test_parameter_ordinals_match_list_index(service, sdk, mocker, catalogue): + catalogue.resolve.return_value = 1 + args = {2: _arg(2, "c", "int"), 0: _arg(0, "a", "int"), 1: _arg(1, "b", "int")} + _attach(service, mocker, _function(args=args)) - service.push_local_function_types_batch({1: 0x1000}, analysis_id=9) + service._push_change(42, 0x1000) - body = sdk.batch_update_function_data_types.call_args.kwargs[ - "batch_update_data_types_input_body" + body = sdk.v3_update_function_signature.call_args.kwargs[ + "update_function_signature_input_body" ] - assert body.functions[0].data_types_version == 5 - - -def test_push_local_types_batch_retries_on_version_conflict(service, sdk, mocker): - service._deci = MagicMock() - service._deci.binary_base_addr = 0 - mocker.patch.object(service, "_build_function_info", return_value=_function_info()) - sdk.list_function_data_types_for_functions.side_effect = [ - MagicMock( - status=True, - data=MagicMock(items=[SimpleNamespace(function_id=1, data_types_version=1)]), - ), - MagicMock( - status=True, - data=MagicMock(items=[SimpleNamespace(function_id=1, data_types_version=2)]), - ), - ] - conflict = MagicMock() - conflict.results = [ - SimpleNamespace(function_id=1, status="version_conflict", data_types_version=None) - ] - ok = MagicMock() - ok.results = [SimpleNamespace(function_id=1, status="updated", data_types_version=3)] - sdk.batch_update_function_data_types.side_effect = [conflict, ok] - - updated = service.push_local_function_types_batch({1: 0x1000}, analysis_id=9) - - assert updated == 1 - assert sdk.batch_update_function_data_types.call_count == 2 + assert [p["ordinal"] for p in body["parameters"]] == [0, 1, 2] + assert [p["name"] for p in body["parameters"]] == ["a", "b", "c"] -def test_resolve_type_builds_valid_function_info_dependency(service, mocker): - from libbs.artifacts import Typedef - from revengai import FunctionDependency, FunctionInfo +def test_unresolved_type_is_sent_as_none(service, sdk, mocker, catalogue): + catalogue.resolve.return_value = None + _attach(service, mocker, _function(args={0: _arg(0, "x", "MysteryType")}, ret=None)) - mocker.patch.object( - svc_mod, "_read_named_type", return_value=Typedef(name="u32", type_="unsigned int") - ) - - dep, referenced = service._resolve_type("u32") - - assert isinstance(dep, FunctionDependency) - assert referenced == ["unsigned int"] - info = FunctionInfo(func_types=None, func_deps=[dep]) - assert info.to_dict()["func_deps"][0] == { - "artifact_type": "Typedef", - "name": "u32", - "type": "unsigned int", - } + service._push_change(42, 0x1000) - -def test_patch_stack_var_type_change_keeps_name(service): - info = _function_info(stack_vars={"-0x20": _sdk_stack_var(-32, "local_20", "char")}) - - changed = service._patch(info, StackVariable(stack_offset=-32, name=None, type_="int", size=4, addr=0x0)) - - assert changed is True - entry = info.func_types.stack_vars["-0x20"] - assert entry.name == "local_20" - assert entry.type == "int" + body = sdk.v3_update_function_signature.call_args.kwargs[ + "update_function_signature_input_body" + ] + assert body["parameters"][0]["data_type_id"] is None + assert body["return_data_type_id"] is None -def test_patch_header_arg_type_change_keeps_name(service): - info = _function_info(args={"0x0": _sdk_arg(0, "oldfile", "char *")}, ret_type="int") - fheader = FunctionHeader( - name=None, addr=0x0, type_=None, args={0: FunctionArgument(offset=0, name=None, type_="wchar_t *", size=8)} +def test_not_found_is_treated_as_nothing_to_edit(service, sdk, mocker, caplog): + sdk.v3_update_function_signature.side_effect = NotFoundException( + status=404, reason="Not Found" ) + _attach(service, mocker, _function()) - changed = service._patch(info, fheader) + service._push_change(42, 0x1000) - assert changed is True - assert info.func_types.header.args["0x0"].name == "oldfile" - assert info.func_types.header.args["0x0"].type == "wchar_t *" + assert sdk.v3_update_function_signature.called -def test_patch_header_identical_is_noop(service): - info = _function_info(args={"0x0": _sdk_arg(0, "a1", "int")}, ret_type="void") - fheader = FunctionHeader( - name="f", addr=0x0, type_="void", args={0: FunctionArgument(offset=0, name="a1", type_="int", size=8)} +def test_ensure_is_given_the_local_type_closure(service, sdk, mocker, catalogue): + node = Struct( + name="Node", + size=8, + members={0: StructMember(name="next", offset=0, type_="Alias", size=8)}, ) - - assert service._patch(info, fheader) is False - - -def test_push_change_fetches_patches_and_pushes(service, sdk, netstore): - netstore.get_analysis_id.return_value = 7 - info = _function_info(stack_vars={"-0x20": _sdk_stack_var(-32, "local_20", "char")}) - fetched = MagicMock() - fetched.status = True - fetched.data.items = [ - MagicMock(function_id=42, data_types=info, data_types_version=3) - ] - sdk.list_function_data_types_for_functions.return_value = fetched - sdk.batch_update_function_data_types.return_value = MagicMock( - results=[MagicMock(status="updated")] + alias = Typedef(name="Alias", type_="int") + named = {"Node": node, "Alias": alias} + mocker.patch.object( + svc_mod, "_read_named_type", side_effect=lambda _d, name: named.get(name) ) + _attach(service, mocker, _function(args={0: _arg(0, "n", "Node *")}, ret="int")) - service._push_change(42, 0x2668, StackVariable(stack_offset=-32, name="n", type_="int", size=4, addr=0x0)) - - sdk.batch_update_function_data_types.assert_called_once() - body = sdk.batch_update_function_data_types.call_args.kwargs[ - "batch_update_data_types_input_body" - ] - item = body.functions[0] - assert item.function_id == 42 - assert item.data_types_version == 3 - - -def test_push_change_skips_push_when_unchanged(service, sdk, netstore): - netstore.get_analysis_id.return_value = 7 - info = _function_info(stack_vars={"-0x20": _sdk_stack_var(-32, "local_20", "char")}) - fetched = MagicMock() - fetched.status = True - fetched.data.items = [ - MagicMock(function_id=42, data_types=info, data_types_version=3) - ] - sdk.list_function_data_types_for_functions.return_value = fetched - - service._push_change(42, 0x2668, StackVariable(stack_offset=-999, name="n", type_="int", size=4, addr=0x0)) - - sdk.batch_update_function_data_types.assert_not_called() + service._push_change(42, 0x1000) + passed = catalogue.ensure.call_args.args[1] + assert set(passed) == {"Node", "Alias"} -def test_push_change_retries_on_version_conflict(service, sdk, netstore): - netstore.get_analysis_id.return_value = 7 - def fresh_response(): - info = _function_info(stack_vars={"-0x20": _sdk_stack_var(-32, "local_20", "char")}) - resp = MagicMock() - resp.status = True - resp.data.items = [MagicMock(function_id=42, data_types=info, data_types_version=3)] - return resp - - sdk.list_function_data_types_for_functions.side_effect = lambda function_ids: fresh_response() - sdk.batch_update_function_data_types.side_effect = [ - MagicMock(results=[MagicMock(status="version_conflict")]), - MagicMock(results=[MagicMock(status="updated")]), - ] - - service._push_change(42, 0x2668, StackVariable(stack_offset=-32, name="n", type_="int", size=4, addr=0x0)) - - assert sdk.batch_update_function_data_types.call_count == 2 - - -def test_push_change_builds_object_when_no_stored_types(service, sdk, netstore, mocker): - netstore.get_analysis_id.return_value = 7 - empty = MagicMock() - empty.status = True - empty.data.items = [] - sdk.list_function_data_types_for_functions.return_value = empty - sdk.batch_update_function_data_types.return_value = MagicMock( - results=[MagicMock(status="updated")] - ) - - func = Function( - addr=0x2668, - size=10, - name="relink", - header=FunctionHeader( - name="relink", - addr=0x2668, - type_="int", - args={0: FunctionArgument(offset=0, name="oldfile", type_="wchar_t *", size=8)}, - ), - stack_vars={-32: StackVariable(stack_offset=-32, name="v2", type_="dev_t", size=8, addr=0x2668)}, - ) - service.attach_decompiler(MagicMock()) - mocker.patch.object(svc_mod, "_read_decompiler_function", return_value=func) - - service._push_change( - 2015112411, - 0x2668, - FunctionHeader(name=None, addr=0x2668, type_=None, args={0: FunctionArgument(offset=0, name=None, type_="wchar_t *", size=8)}), +def test_stack_variable_types_reach_the_catalogue(service, sdk, mocker, catalogue): + local = Struct(name="Local", size=4, members={}) + mocker.patch.object( + svc_mod, "_read_named_type", side_effect=lambda _d, name: {"Local": local}.get(name) ) + stack_vars = { + 0: StackVariable(stack_offset=0, name="v", type_="Local", size=4, addr=0x1000) + } + _attach(service, mocker, _function(stack_vars=stack_vars, ret=None)) - sdk.batch_update_function_data_types.assert_called_once() - body = sdk.batch_update_function_data_types.call_args.kwargs[ - "batch_update_data_types_input_body" - ] - item = body.functions[0] - assert item.function_id == 2015112411 - assert item.data_types_version == 0 - assert item.data_types["func_types"]["header"]["args"]["0x0"]["name"] == "oldfile" - assert item.data_types["func_types"]["header"]["args"]["0x0"]["type"] == "wchar_t *" - assert item.data_types["func_types"]["stack_vars"]["-0x20"]["name"] == "v2" - - -def test_push_change_skips_build_when_no_decompiler(service, sdk, netstore): - netstore.get_analysis_id.return_value = 7 - empty = MagicMock() - empty.status = True - empty.data.items = [] - sdk.list_function_data_types_for_functions.return_value = empty - - service._push_change(42, 0x2668, StackVariable(stack_offset=-32, name="n", type_="int", size=4, addr=0x0)) - - sdk.batch_update_function_data_types.assert_not_called() + service._push_change(42, 0x1000) + passed = catalogue.ensure.call_args.args[1] + assert "Local" in passed -def test_collect_func_deps_resolves_typedef_chain(service, mocker): - from libbs.artifacts import Typedef - service.attach_decompiler(MagicMock()) - func_type = _function_info(args={"0x0": _sdk_arg(0, "d", "dev_t")}, ret_type="int").func_types +def test_type_closure_is_capped(service, sdk, mocker, catalogue): + def endless(_deci, name): + index = int(name[1:]) if name[1:].isdigit() else 0 + return Typedef(name=name, type_=f"T{index + 1}") - chain = { - "dev_t": Typedef(name="dev_t", type_="__dev_t"), - "__dev_t": Typedef(name="__dev_t", type_="unsigned long"), - } - mocker.patch.object(svc_mod, "_read_named_type", side_effect=lambda deci, name: chain.get(name)) + mocker.patch.object(svc_mod, "_read_named_type", side_effect=endless) + _attach(service, mocker, _function(args={0: _arg(0, "x", "T0")}, ret=None)) - deps = service._collect_func_deps(func_type) + service._push_change(42, 0x1000) - assert sorted(d.name for d in deps) == ["__dev_t", "dev_t"] - assert all(d.artifact_type == "Typedef" for d in deps) + passed = catalogue.ensure.call_args.args[1] + assert len(passed) == svc_mod.MAX_TYPE_DEPENDENCIES -def test_collect_func_deps_resolves_struct_members(service, mocker): - from libbs.artifacts import Struct, StructMember, Typedef +@pytest.mark.parametrize( + "raw,expected", + [ + ("Foo *", "Foo"), + ("const struct Bar *", "Bar"), + ("unsigned int", "int"), + ("char[16]", "char"), + ("", None), + (None, None), + ], +) +def test_base_type_name(raw, expected): + assert VariableSyncService._base_type_name(raw) == expected - service.attach_decompiler(MagicMock()) - func_type = _function_info(args={"0x0": _sdk_arg(0, "p", "mystruct *")}, ret_type="int").func_types - types = { - "mystruct": Struct( - name="mystruct", size=8, members={0: StructMember(name="x", offset=0, type_="myint", size=4)} - ), - "myint": Typedef(name="myint", type_="int"), - } - mocker.patch.object(svc_mod, "_read_named_type", side_effect=lambda deci, name: types.get(name)) +def test_push_local_types_batch_noop_without_targets_or_analysis(service, sdk): + assert service.push_local_function_types_batch({}, ANALYSIS) == 0 + assert service.push_local_function_types_batch({1: 0x1000}, None) == 0 + sdk.v3_update_function_signature.assert_not_called() - deps = service._collect_func_deps(func_type) - by_name = {d.name: d for d in deps} - assert set(by_name) == {"mystruct", "myint"} - assert by_name["mystruct"].members["0x0"]["type"] == "myint" +def test_push_local_types_batch_counts_successful_puts(service, sdk, mocker, catalogue): + catalogue.resolve.return_value = 1 + deci = MagicMock() + deci.binary_base_addr = 0 + service.attach_decompiler(deci) + mocker.patch.object(svc_mod, "_read_decompiler_function", return_value=_function()) + assert service.push_local_function_types_batch({1: 0x1000, 2: 0x2000}, ANALYSIS) == 2 + assert sdk.v3_update_function_signature.call_count == 2 -def test_push_change_no_analysis_id_does_nothing(service, sdk, netstore): - netstore.get_analysis_id.return_value = None - service._push_change(42, 0x2668, StackVariable(stack_offset=-32, name="n", type_="int", size=4, addr=0x0)) +def test_push_local_types_batch_skips_functions_without_a_stored_signature( + service, sdk, mocker, catalogue +): + sdk.v3_update_function_signature.side_effect = [ + None, + NotFoundException(status=404, reason="Not Found"), + ] + deci = MagicMock() + deci.binary_base_addr = 0 + service.attach_decompiler(deci) + mocker.patch.object(svc_mod, "_read_decompiler_function", return_value=_function()) - sdk.list_function_data_types_for_functions.assert_not_called() + assert service.push_local_function_types_batch({1: 0x1000, 2: 0x2000}, ANALYSIS) == 1 def test_enqueue_change_debounces_rapid_duplicates(service, mocker): - mocker.patch.object(service, "_start_worker_if_needed") - svar = StackVariable(stack_offset=-32, name="a", type_="int", size=4, addr=0x0) + mocker.patch.object(VariableSyncService, "_start_worker_if_needed") + header = FunctionHeader(name="f", addr=0x1000) - service.enqueue_change(42, 0x2668, svar) - service.enqueue_change(42, 0x2668, svar) + service.enqueue_change(1, 0x1000, header) + service.enqueue_change(1, 0x1000, header) assert service._q.qsize() == 1 From 2357dad167e35038fce09627447112e778070e34 Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:22:22 +0100 Subject: [PATCH 04/12] feat: AI decompilation token resolution --- .../app/coordinators/ai_decomp_coordinator.py | 6 +- .../app/coordinators/ai_decomp_render.py | 99 +++++++------ .../services/ai_decomp/ai_decomp_service.py | 51 +++---- tests/idalib/ai_decomp/test_smoke.py | 24 ++- tests/ui/ai_decomp_ui_runner.py | 33 +++-- tests/unit/ai_decomp/test_render.py | 140 ++++++++++++------ tests/unit/ai_decomp/test_sdk_schemas.py | 117 ++++++++++----- tests/unit/ai_decomp/test_service.py | 67 ++++++--- 8 files changed, 346 insertions(+), 191 deletions(-) diff --git a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py index 7eb7aea..6868e3d 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py +++ b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py @@ -7,7 +7,7 @@ from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData from revengai.models.summary_data import SummaryData -from revengai.models.tokenised_data import TokenisedData +from revengai.models.get_tokens_response import GetTokensResponse from revengai.models.workflow_progress import WorkflowProgress from reai_toolkit.app.app import App @@ -43,7 +43,7 @@ def __init__( self._current_decomp: DecompilationData | None = None self._current_summary: SummaryData | None = None self._current_comments: CommentsData | None = None - self._current_tokenised: TokenisedData | None = None + self._current_tokenised: GetTokensResponse | None = None self._baseline: RenderModel | None = None def enable_function_tracking(self) -> None: @@ -192,7 +192,7 @@ def _on_comments_complete( self._rerender() def _on_tokenised_complete( - self, ea: int, response: GenericApiReturn[TokenisedData] + self, ea: int, response: GenericApiReturn[GetTokensResponse] ) -> None: if ea != self._current_func_vaddr: return diff --git a/reai_toolkit/app/coordinators/ai_decomp_render.py b/reai_toolkit/app/coordinators/ai_decomp_render.py index 67513fc..8650483 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_render.py +++ b/reai_toolkit/app/coordinators/ai_decomp_render.py @@ -5,22 +5,17 @@ from typing import TYPE_CHECKING, Optional if TYPE_CHECKING: - from revengai.models.ai_decomp_function_mapping import AIDecompFunctionMapping from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData + from revengai.models.get_tokens_response import GetTokensResponse from revengai.models.progress_message import ProgressMessage + from revengai.models.rendered_token import RenderedToken from revengai.models.summary_data import SummaryData - from revengai.models.tokenised_data import TokenisedData from revengai.models.workflow_progress import WorkflowProgress _IDENT_RE = re.compile(r"[A-Za-z_]\w*") -_VARIABLE_CATEGORIES = ( - "unmatched_vars", - "unmatched_global_vars", - "unmatched_external_vars", -) -_TYPE_CATEGORIES = ("unmatched_custom_types", "unmatched_enums") +_RENAMED_ELSEWHERE_FIELDS = ("data_type_id", "function_id", "imported_function_id") @dataclass @@ -170,53 +165,65 @@ def index_of_identifier(line: str, word: str) -> int: return idents.index(word) if word in idents else -1 -def resolve_token( - tokenised: "TokenisedData", +def effective_values(tokens: "GetTokensResponse") -> dict[str, str]: + rendered = tokens.placeholder_to_rendered_token or {} + overrides = tokens.placeholder_to_user_override or {} + values: dict[str, str] = {} + for placeholder, token in rendered.items(): + override = overrides.get(placeholder) + values[placeholder] = override.value if override is not None else token.value + return values + + +def names_token(ident: str, rendered_value: str) -> bool: + return ident == rendered_value or ident in _IDENT_RE.findall(rendered_value) + + +def _unit_re(placeholders) -> re.Pattern: + parts = [re.escape(p) for p in sorted(placeholders, key=len, reverse=True)] + parts.append(_IDENT_RE.pattern) + return re.compile("|".join(parts)) + + +def find_token( + tokens: "GetTokensResponse", source_index: int, ident_index: int, old_ident: str, -) -> Optional[tuple[str, str]]: - mapping = tokenised.function_mapping - if mapping is None: +) -> Optional[tuple[str, "RenderedToken"]]: + rendered = tokens.placeholder_to_rendered_token or {} + if not rendered: return None + values = effective_values(tokens) - tok_lines = (tokenised.tokenised_decompilation or "").split("\n") + tok_lines = (tokens.ai_decomp or "").split("\n") if 0 <= source_index < len(tok_lines): - tok_idents = _IDENT_RE.findall(tok_lines[source_index]) - if 0 <= ident_index < len(tok_idents): - candidate = tok_idents[ident_index] - cat, eff = _category_of_token(mapping, candidate) - if cat is not None and eff == old_ident: - return candidate, cat - - matches = [ - (token, cat) - for token, rv, cat in _iter_category_tokens(mapping) - if _effective_value(mapping, token, rv) == old_ident - ] + units = _unit_re(rendered).findall(tok_lines[source_index]) + if 0 <= ident_index < len(units): + unit = units[ident_index] + if unit in rendered and names_token(old_ident, values[unit]): + return unit, rendered[unit] + + matches = [p for p, value in values.items() if names_token(old_ident, value)] if len(matches) == 1: - return matches[0] + return matches[0], rendered[matches[0]] return None -def _iter_category_tokens(mapping: "AIDecompFunctionMapping"): - for cat in _VARIABLE_CATEGORIES: - for token, rv in (getattr(mapping, cat, None) or {}).items(): - yield token, rv, "variable" - for cat in _TYPE_CATEGORIES: - for token, rv in (getattr(mapping, cat, None) or {}).items(): - yield token, rv, "type" +def is_renameable(token: "RenderedToken") -> bool: + return all(getattr(token, name, None) is None for name in _RENAMED_ELSEWHERE_FIELDS) -def _category_of_token( - mapping: "AIDecompFunctionMapping", token: str -) -> tuple[Optional[str], Optional[str]]: - for t, rv, cat in _iter_category_tokens(mapping): - if t == token: - return cat, _effective_value(mapping, t, rv) - return None, None - - -def _effective_value(mapping: "AIDecompFunctionMapping", token: str, replacement) -> str: - overrides = mapping.user_override_mappings or {} - return overrides.get(token, replacement.value) +def resolve_token( + tokens: "GetTokensResponse", + source_index: int, + ident_index: int, + old_ident: str, +) -> Optional[tuple[str, str]]: + found = find_token(tokens, source_index, ident_index, old_ident) + if found is None: + return None + placeholder, token = found + if not is_renameable(token): + return None + return placeholder, token.kind diff --git a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py index c53a227..df961be 100644 --- a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py +++ b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py @@ -11,10 +11,11 @@ from revengai.models.ai_decompilation_rating import AiDecompilationRating from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData +from revengai.models.get_tokens_response import GetTokensResponse from revengai.models.patch_comment_body import PatchCommentBody from revengai.models.summary_data import SummaryData from revengai.models.task_status import TaskStatus -from revengai.models.tokenised_data import TokenisedData +from revengai.models.token import Token from revengai.models.upsert_ai_decomplation_rating_request import ( UpsertAiDecomplationRatingRequest, ) @@ -37,7 +38,7 @@ def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration) self._decomp_cache: dict[int, DecompilationData] = {} self._summary_cache: dict[int, SummaryData] = {} self._comments_cache: dict[int, CommentsData] = {} - self._tokenised_cache: dict[int, TokenisedData] = {} + self._tokenised_cache: dict[int, GetTokensResponse] = {} self._inflight: dict[int, threading.Event] = {} self._inflight_lock = threading.Lock() @@ -89,7 +90,7 @@ def start_ai_decomp_task( on_decomp: Callable[[GenericApiReturn[DecompilationData]], None], on_summary: Callable[[GenericApiReturn[SummaryData]], None], on_comments: Callable[[GenericApiReturn[CommentsData]], None], - on_tokenised: Callable[[GenericApiReturn[TokenisedData]], None] | None = None, + on_tokenised: Callable[[GenericApiReturn[GetTokensResponse]], None] | None = None, on_progress: Callable[[WorkflowProgress], None] | None = None, ) -> None: function_id = self._get_function_id(start_ea=ea) @@ -124,7 +125,7 @@ def apply_overrides( ea: int, overrides: dict[str, str], on_decomp: Callable[[GenericApiReturn[DecompilationData]], None], - on_tokenised: Callable[[GenericApiReturn[TokenisedData]], None], + on_tokenised: Callable[[GenericApiReturn[GetTokensResponse]], None], ) -> None: function_id = self._get_function_id(start_ea=ea) if function_id is None: @@ -240,14 +241,17 @@ def _run_apply_overrides( function_id: int, overrides: dict[str, str], on_decomp: Callable[[GenericApiReturn[DecompilationData]], None], - on_tokenised: Callable[[GenericApiReturn[TokenisedData]], None], + on_tokenised: Callable[[GenericApiReturn[GetTokensResponse]], None], ) -> None: try: with self.yield_api_client(sdk_config=self.sdk_config) as api_client: - FunctionsAIDecompilationApi(api_client).upsert_ai_decompilation_overrides( + FunctionsAIDecompilationApi(api_client).v3_upsert_ai_decompilation_overrides( function_id=function_id, upsert_overrides_input_body=UpsertOverridesInputBody( - overrides=overrides + overrides={ + placeholder: Token(value=value) + for placeholder, value in overrides.items() + } ), ) except ApiException as e: @@ -280,13 +284,9 @@ def _run_apply_overrides( ) tokenised, _terr = self._fetch_tokenised(function_id) - if ( - tokenised is not None - and str(tokenised.status) == TaskStatus.COMPLETED.value - and tokenised.function_mapping is not None - ): + if tokenised is not None and _tokens_ready(tokenised): self._tokenised_cache[function_id] = tokenised - on_tokenised(GenericApiReturn[TokenisedData](success=True, data=tokenised)) + on_tokenised(GenericApiReturn[GetTokensResponse](success=True, data=tokenised)) def _run_set_comment( self, @@ -358,7 +358,7 @@ def _run_task( on_decomp: Callable[[GenericApiReturn[DecompilationData]], None], on_summary: Callable[[GenericApiReturn[SummaryData]], None], on_comments: Callable[[GenericApiReturn[CommentsData]], None], - on_tokenised: Callable[[GenericApiReturn[TokenisedData]], None] | None = None, + on_tokenised: Callable[[GenericApiReturn[GetTokensResponse]], None] | None = None, on_progress: Callable[[WorkflowProgress], None] | None = None, ) -> None: try: @@ -781,12 +781,12 @@ def _run_comments_phase( def _fetch_tokenised( self, function_id: int - ) -> tuple[TokenisedData | None, str | None]: + ) -> tuple[GetTokensResponse | None, str | None]: try: with self.yield_api_client(sdk_config=self.sdk_config) as api_client: tokenised = FunctionsAIDecompilationApi( api_client - ).get_ai_decompilation_tokenised(function_id=function_id) + ).v3_get_ai_decompilation_tokens(function_id=function_id) return tokenised, None except ApiException as e: return None, _format_api_error(e) @@ -797,14 +797,14 @@ def _run_tokenised_phase( self, function_id: int, stop_event: threading.Event, - on_tokenised: Callable[[GenericApiReturn[TokenisedData]], None], + on_tokenised: Callable[[GenericApiReturn[GetTokensResponse]], None], ) -> None: cached = self._tokenised_cache.get(function_id) if cached is not None: self._safe_dispatch( stop_event, on_tokenised, - GenericApiReturn[TokenisedData](success=True, data=cached), + GenericApiReturn[GetTokensResponse](success=True, data=cached), ) return @@ -813,25 +813,22 @@ def _run_tokenised_phase( self._safe_dispatch( stop_event, on_tokenised, - GenericApiReturn[TokenisedData](success=False, error_message=err), + GenericApiReturn[GetTokensResponse](success=False, error_message=err), ) return - if ( - str(tokenised.status) == TaskStatus.COMPLETED.value - and tokenised.function_mapping is not None - ): + if _tokens_ready(tokenised): self._tokenised_cache[function_id] = tokenised self._safe_dispatch( stop_event, on_tokenised, - GenericApiReturn[TokenisedData](success=True, data=tokenised), + GenericApiReturn[GetTokensResponse](success=True, data=tokenised), ) else: self._safe_dispatch( stop_event, on_tokenised, - GenericApiReturn[TokenisedData]( + GenericApiReturn[GetTokensResponse]( success=False, error_message="Tokenised data not ready." ), ) @@ -883,6 +880,10 @@ def _poll_workflow( return False, None +def _tokens_ready(tokens: GetTokensResponse) -> bool: + return bool(tokens.ai_decomp) and bool(tokens.placeholder_to_rendered_token) + + def _format_api_error(e: ApiException) -> str: error_response = parse_exception(e) if error_response and error_response.errors: diff --git a/tests/idalib/ai_decomp/test_smoke.py b/tests/idalib/ai_decomp/test_smoke.py index 37610e8..cb860d9 100644 --- a/tests/idalib/ai_decomp/test_smoke.py +++ b/tests/idalib/ai_decomp/test_smoke.py @@ -9,7 +9,8 @@ from revengai.models.inline_comment import InlineComment from revengai.models.summary_data import SummaryData from revengai.models.task_status import TaskStatus -from revengai.models.tokenised_data import TokenisedData +from revengai.models.get_tokens_response import GetTokensResponse +from revengai.models.rendered_token import RenderedToken from reai_toolkit.app.services.ai_decomp import ai_decomp_service as svc_mod from reai_toolkit.app.services.ai_decomp.ai_decomp_service import AiDecompService @@ -75,11 +76,20 @@ def test_mutations_and_refresh_under_idalib(loaded_binary, mocker): api_inst.get_ai_decompilation.return_value = DecompilationData.model_construct( status=TaskStatus.COMPLETED.value, decompilation="int sub(int v1) { return v1; }" ) - api_inst.get_ai_decompilation_tokenised.return_value = TokenisedData.model_construct( - status=TaskStatus.COMPLETED.value, - tokenised_decompilation="int @@F@@(int @@V@@) { return @@V@@; }", - predicted_function_name="sub", - function_mapping=MagicMock(), + api_inst.v3_get_ai_decompilation_tokens.return_value = GetTokensResponse.model_construct( + ai_decomp="int @@F@@(int @@V@@) { return @@V@@; }", + analysis_id=1, + placeholder_to_rendered_token={ + "@@V@@": RenderedToken.model_construct( + value="v1", + kind="param", + vaddr=None, + data_type_id=None, + function_id=None, + imported_function_id=None, + ) + }, + placeholder_to_user_override={}, ) api_inst.patch_ai_decompilation_inline_comment.return_value = ( CommentsData.model_construct( @@ -103,7 +113,7 @@ def test_mutations_and_refresh_under_idalib(loaded_binary, mocker): ) _wait_mock(on_decomp) _wait_mock(on_tok) - api_inst.upsert_ai_decompilation_overrides.assert_called_once() + api_inst.v3_upsert_ai_decompilation_overrides.assert_called_once() assert on_decomp.call_args[0][0].success is True assert service._tokenised_cache[7] is not None diff --git a/tests/ui/ai_decomp_ui_runner.py b/tests/ui/ai_decomp_ui_runner.py index 6cc5b43..4a92feb 100644 --- a/tests/ui/ai_decomp_ui_runner.py +++ b/tests/ui/ai_decomp_ui_runner.py @@ -34,12 +34,11 @@ def _run(report: dict) -> None: import ida_kernwin import idautils - from revengai.models.ai_decomp_function_mapping import AIDecompFunctionMapping from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData + from revengai.models.get_tokens_response import GetTokensResponse from revengai.models.inline_comment import InlineComment - from revengai.models.replacement_value import ReplacementValue - from revengai.models.tokenised_data import TokenisedData + from revengai.models.rendered_token import RenderedToken from reai_toolkit.app.components.tabs.ai_decomp_tab import AIDecompView from reai_toolkit.app.coordinators.ai_decomp_coordinator import AiDecompCoordinator @@ -53,18 +52,24 @@ def _run(report: dict) -> None: ea = next(iter(idautils.Functions()), 0x1000) decomp = DecompilationData.model_construct(status="COMPLETED", decompilation=CODE) - mapping = AIDecompFunctionMapping.model_construct( - unmatched_vars={ - "@@A@@": ReplacementValue.model_construct(value="a1"), - "@@V@@": ReplacementValue.model_construct(value="v5"), + def _rt(value, kind): + return RenderedToken.model_construct( + value=value, + kind=kind, + vaddr=None, + data_type_id=None, + function_id=None, + imported_function_id=None, + ) + + tokenised = GetTokensResponse.model_construct( + ai_decomp=TOK, + analysis_id=1, + placeholder_to_rendered_token={ + "@@A@@": _rt("a1", "param"), + "@@V@@": _rt("v5", "local"), }, - user_override_mappings={}, - ) - tokenised = TokenisedData.model_construct( - status="COMPLETED", - tokenised_decompilation=TOK, - predicted_function_name="f", - function_mapping=mapping, + placeholder_to_user_override={}, ) service = MagicMock() diff --git a/tests/unit/ai_decomp/test_render.py b/tests/unit/ai_decomp/test_render.py index d13ab3a..f5c2a33 100644 --- a/tests/unit/ai_decomp/test_render.py +++ b/tests/unit/ai_decomp/test_render.py @@ -1,15 +1,19 @@ -from revengai.models.ai_decomp_function_mapping import AIDecompFunctionMapping +import pytest from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData +from revengai.models.get_tokens_response import GetTokensResponse from revengai.models.inline_comment import InlineComment from revengai.models.progress_message import ProgressMessage -from revengai.models.replacement_value import ReplacementValue +from revengai.models.rendered_token import RenderedToken from revengai.models.summary_data import SummaryData -from revengai.models.tokenised_data import TokenisedData +from revengai.models.token import Token from revengai.models.workflow_progress import WorkflowProgress from reai_toolkit.app.coordinators.ai_decomp_render import ( + effective_values, + find_token, index_of_identifier, + names_token, render_progress, render_view, render_view_with_map, @@ -34,38 +38,30 @@ def _comments(pairs): return CommentsData.model_construct(inline_comments=items, task_status="COMPLETED") -def _mapping(**cats): - base = dict( - fields={}, - inverse_function_map={}, - inverse_string_map={}, - unmatched_custom_function_pointers={}, - unmatched_custom_types={}, - unmatched_enums={}, - unmatched_external_vars={}, - unmatched_functions={}, - unmatched_global_vars={}, - unmatched_go_to_labels={}, - unmatched_strings={}, - unmatched_variadic_lists={}, - unmatched_vars={}, - user_override_mappings={}, +def _rt(value, kind="local", **ids): + return RenderedToken.model_construct( + value=value, + kind=kind, + vaddr=None, + data_type_id=ids.get("data_type_id"), + function_id=ids.get("function_id"), + imported_function_id=ids.get("imported_function_id"), ) - base.update(cats) - return AIDecompFunctionMapping.model_construct(**base) -def _rv(value): - return ReplacementValue.model_construct(value=value) +def _tokens(tok=TOK, rendered=None, overrides=None): + return GetTokensResponse.model_construct( + ai_decomp=tok, + analysis_id=1, + placeholder_to_rendered_token=rendered if rendered is not None else {}, + placeholder_to_user_override={ + placeholder: Token.model_construct(value=value) + for placeholder, value in (overrides or {}).items() + }, + ) -def _tokenised(tok=TOK, mapping=None): - return TokenisedData.model_construct( - status="COMPLETED", - tokenised_decompilation=tok, - predicted_function_name="f", - function_mapping=mapping if mapping is not None else _mapping(), - ) +_VARS = {"@@V_v5@@": _rt("v5"), "@@V_a1@@": _rt("a1", kind="param")} def test_render_matches_legacy_and_builds_model(): @@ -128,33 +124,85 @@ def test_index_of_identifier(): def test_resolve_token_positional_variable(): - tokd = _tokenised(mapping=_mapping(unmatched_vars={"@@V_v5@@": _rv("v5"), "@@V_a1@@": _rv("a1")})) - assert resolve_token(tokd, 1, 1, "v5") == ("@@V_v5@@", "variable") + assert resolve_token(_tokens(rendered=_VARS), 1, 1, "v5") == ("@@V_v5@@", "local") -def test_resolve_token_type_category(): - tok = "@@T_S@@ *x = 0;" - tokd = _tokenised(tok=tok, mapping=_mapping(unmatched_custom_types={"@@T_S@@": _rv("Foo")})) - assert resolve_token(tokd, 0, 0, "Foo") == ("@@T_S@@", "type") +def test_resolve_token_picks_the_placeholder_at_that_position_not_a_same_valued_twin(): + rendered = {"@@V_a@@": _rt("dup"), "@@V_b@@": _rt("dup")} + tokens = _tokens(tok="int @@V_a@@ = @@V_b@@;", rendered=rendered) + + assert resolve_token(tokens, 0, 2, "dup") == ("@@V_b@@", "local") def test_resolve_token_honours_user_override(): - mapping = _mapping( - unmatched_vars={"@@V_v5@@": _rv("v5")}, - user_override_mappings={"@@V_v5@@": "tmp"}, - ) - tokd = _tokenised(mapping=mapping) - assert resolve_token(tokd, 1, 1, "tmp") == ("@@V_v5@@", "variable") + tokens = _tokens(rendered=_VARS, overrides={"@@V_v5@@": "tmp"}) + + assert resolve_token(tokens, 1, 1, "tmp") == ("@@V_v5@@", "local") + assert resolve_token(tokens, 1, 1, "v5") is None + + +def test_effective_values_merge_overrides_over_rendered(): + tokens = _tokens(rendered=_VARS, overrides={"@@V_v5@@": "tmp"}) + + assert effective_values(tokens) == {"@@V_v5@@": "tmp", "@@V_a1@@": "a1"} def test_resolve_token_unknown_returns_none(): - tokd = _tokenised(mapping=_mapping(unmatched_vars={"@@V_v5@@": _rv("v5")})) - assert resolve_token(tokd, 1, 1, "not_a_var") is None + assert resolve_token(_tokens(rendered=_VARS), 1, 1, "not_a_var") is None def test_resolve_token_value_fallback_when_line_unaligned(): - tokd = _tokenised(tok="", mapping=_mapping(unmatched_vars={"@@V_v5@@": _rv("v5")})) - assert resolve_token(tokd, 1, 1, "v5") == ("@@V_v5@@", "variable") + tokens = _tokens(tok="", rendered={"@@V_v5@@": _rt("v5")}) + + assert resolve_token(tokens, 1, 1, "v5") == ("@@V_v5@@", "local") + + +def test_resolve_token_ambiguous_value_fallback_resolves_to_nothing(): + rendered = {"@@V_a@@": _rt("dup"), "@@V_b@@": _rt("dup")} + + assert resolve_token(_tokens(tok="", rendered=rendered), 0, 0, "dup") is None + + +def test_resolve_token_empty_token_map_returns_none(): + assert resolve_token(_tokens(rendered={}), 1, 1, "v5") is None + + +@pytest.mark.parametrize( + "value,ident", + [("lang_start<()>", "lang_start"), ("Foo::bar", "bar"), ("Foo::bar", "Foo")], +) +def test_identifier_inside_a_qualified_rendered_value_still_resolves(value, ident): + tokens = _tokens(tok="@@F@@();", rendered={"@@F@@": _rt(value)}) + + assert names_token(ident, value) + assert resolve_token(tokens, 0, 0, ident) == ("@@F@@", "local") + + +@pytest.mark.parametrize( + "field", ["data_type_id", "function_id", "imported_function_id"] +) +def test_tokens_renamed_elsewhere_are_declined(field): + rendered = {"@@X@@": _rt("Foo", kind="type", **{field: 12})} + + assert resolve_token(_tokens(tok="@@X@@ *x;", rendered=rendered), 0, 0, "Foo") is None + + +@pytest.mark.parametrize( + "field", ["data_type_id", "function_id", "imported_function_id"] +) +def test_id_zero_is_a_real_id_and_still_declines(field): + rendered = {"@@X@@": _rt("Foo", kind="type", **{field: 0})} + + assert resolve_token(_tokens(tok="@@X@@ *x;", rendered=rendered), 0, 0, "Foo") is None + + +def test_find_token_returns_tokens_renamed_elsewhere_even_though_resolve_declines(): + rendered = {"@@X@@": _rt("Foo", kind="type", data_type_id=12)} + tokens = _tokens(tok="@@X@@ *x;", rendered=rendered) + + placeholder, token = find_token(tokens, 0, 0, "Foo") + assert placeholder == "@@X@@" + assert token.data_type_id == 12 def _pm(text, level="INFO", step="DECOMPILING", timestamp=None): diff --git a/tests/unit/ai_decomp/test_sdk_schemas.py b/tests/unit/ai_decomp/test_sdk_schemas.py index e3083f9..d70cb6f 100644 --- a/tests/unit/ai_decomp/test_sdk_schemas.py +++ b/tests/unit/ai_decomp/test_sdk_schemas.py @@ -2,32 +2,96 @@ Schema-shape assertions: fail loudly when the revengai SDK shape drifts. """ -from revengai.models.ai_decomp_function_mapping import AIDecompFunctionMapping +import inspect + +from revengai import FunctionsAIDecompilationApi from revengai.models.comments_data import CommentsData from revengai.models.create_ai_decomp_output_body import CreateAIDecompOutputBody +from revengai.models.get_tokens_response import GetTokensResponse from revengai.models.inline_comment import InlineComment from revengai.models.patch_comment_body import PatchCommentBody -from revengai.models.replacement_value import ReplacementValue +from revengai.models.rendered_token import RenderedToken +from revengai.models.summary_data import SummaryData from revengai.models.task_status import TaskStatus -from revengai.models.tokenised_data import TokenisedData +from revengai.models.token import Token from revengai.models.upsert_overrides_data import UpsertOverridesData from revengai.models.upsert_overrides_input_body import UpsertOverridesInputBody from revengai.models.workflow_progress import WorkflowProgress +def test_ai_decomp_api_exposes_the_methods_the_plugin_calls(): + for method in ( + "create_ai_decompilation", + "get_ai_decompilation", + "get_ai_decompilation_status", + "get_ai_decompilation_summary", + "get_ai_decompilation_summary_status", + "regenerate_ai_decompilation_summary", + "get_ai_decompilation_inline_comments", + "get_ai_decompilation_inline_comments_status", + "regenerate_ai_decompilation_inline_comments", + "patch_ai_decompilation_inline_comment", + "delete_ai_decompilation_inline_comment", + "upsert_ai_decompilation_rating", + "v3_get_ai_decompilation_tokens", + "v3_upsert_ai_decompilation_overrides", + ): + assert callable(getattr(FunctionsAIDecompilationApi, method)) + + +def test_v3_token_methods_accept_plugin_kwargs(): + tokens = inspect.signature( + FunctionsAIDecompilationApi.v3_get_ai_decompilation_tokens + ).parameters + assert "function_id" in tokens + + overrides = inspect.signature( + FunctionsAIDecompilationApi.v3_upsert_ai_decompilation_overrides + ).parameters + assert {"function_id", "upsert_overrides_input_body"} <= set(overrides) + + def test_workflow_progress_has_expected_fields(): fields = set(WorkflowProgress.model_fields.keys()) assert {"status", "step", "step_index", "steps_total", "messages"} <= fields -def test_tokenised_data_has_expected_fields(): - fields = set(TokenisedData.model_fields.keys()) +def test_get_tokens_response_carries_source_and_both_token_maps(): assert { - "status", - "tokenised_decompilation", - "predicted_function_name", - "function_mapping", - } <= fields + "ai_decomp", + "analysis_id", + "placeholder_to_rendered_token", + "placeholder_to_user_override", + } <= set(GetTokensResponse.model_fields) + + +def test_get_tokens_response_has_no_status_so_readiness_is_the_empty_source(): + assert "status" not in GetTokensResponse.model_fields + + +def test_rendered_token_carries_value_kind_and_the_renamed_elsewhere_ids(): + assert { + "value", + "kind", + "data_type_id", + "function_id", + "imported_function_id", + } <= set(RenderedToken.model_fields) + + +def test_rendered_token_kinds_cover_what_the_panel_offers_to_rename(): + tokens = GetTokensResponse.from_json( + '{"ai_decomp": "@@V@@;", "analysis_id": 1,' + ' "placeholder_to_rendered_token": {"@@V@@": {"kind": "local", "value": "v5"}},' + ' "placeholder_to_user_override": {}}' + ) + rendered = tokens.placeholder_to_rendered_token["@@V@@"] + assert rendered.value == "v5" + assert rendered.data_type_id is None + + +def test_predicted_function_name_now_rides_on_the_summary(): + assert "predicted_function_name" in SummaryData.model_fields def test_create_output_body_has_status(): @@ -46,15 +110,18 @@ def test_workflow_progress_accepts_uppercase_status_values(): ) -def test_upsert_overrides_input_body_round_trips(): - body = UpsertOverridesInputBody.from_json('{"overrides": {"@@V@@": "buf", "@@T@@": ""}}') - assert body.overrides == {"@@V@@": "buf", "@@T@@": ""} - assert "overrides" in UpsertOverridesInputBody.model_fields +def test_overrides_are_posted_as_token_objects_not_bare_strings(): + body = UpsertOverridesInputBody.from_json( + '{"overrides": {"@@V@@": {"value": "buf"}}}' + ) + assert body.overrides == {"@@V@@": Token(value="buf")} -def test_upsert_overrides_data_exposes_merged_map(): - data = UpsertOverridesData.from_json('{"user_override_mappings": {"@@V@@": "buf"}}') - assert data.user_override_mappings == {"@@V@@": "buf"} +def test_upsert_overrides_data_returns_the_override_map(): + data = UpsertOverridesData.from_json( + '{"placeholder_to_user_override": {"@@V@@": {"value": "buf"}}}' + ) + assert data.placeholder_to_user_override["@@V@@"].value == "buf" def test_patch_comment_body_round_trips(): @@ -67,19 +134,3 @@ def test_patch_comment_body_round_trips(): def test_inline_comment_and_comments_data_shapes(): assert {"comment", "line"} <= set(InlineComment.model_fields.keys()) assert {"inline_comments", "task_status"} <= set(CommentsData.model_fields.keys()) - - -def test_replacement_value_carries_value(): - assert ReplacementValue.from_json('{"value": "count"}').value == "count" - - -def test_function_mapping_has_variable_and_type_categories(): - fields = set(AIDecompFunctionMapping.model_fields.keys()) - assert { - "unmatched_vars", - "unmatched_global_vars", - "unmatched_external_vars", - "unmatched_custom_types", - "unmatched_enums", - "user_override_mappings", - } <= fields diff --git a/tests/unit/ai_decomp/test_service.py b/tests/unit/ai_decomp/test_service.py index c25a210..9d8212b 100644 --- a/tests/unit/ai_decomp/test_service.py +++ b/tests/unit/ai_decomp/test_service.py @@ -10,10 +10,11 @@ from revengai.models.inline_comment import InlineComment from revengai.models.summary_data import SummaryData from revengai.models.task_status import TaskStatus -from revengai.models.tokenised_data import TokenisedData +from revengai.models.get_tokens_response import GetTokensResponse +from revengai.models.rendered_token import RenderedToken +from revengai.models.token import Token from revengai.models.workflow_progress import WorkflowProgress -from reai_toolkit.app.core.shared_schema import GenericApiReturn from reai_toolkit.app.services.ai_decomp import ai_decomp_service as svc_mod from reai_toolkit.app.services.ai_decomp.ai_decomp_service import AiDecompService @@ -72,12 +73,21 @@ def _comments(items=None, status=TaskStatus.COMPLETED.value) -> CommentsData: return CommentsData.model_construct(inline_comments=items or [], task_status=status) -def _tokd(status=TaskStatus.COMPLETED.value) -> TokenisedData: - return TokenisedData.model_construct( - status=status, - tokenised_decompilation="int @@F@@(void) {}", - predicted_function_name="f", - function_mapping=MagicMock(), +def _tokd(source="int @@F@@(void) {}") -> GetTokensResponse: + return GetTokensResponse.model_construct( + ai_decomp=source, + analysis_id=1, + placeholder_to_rendered_token={ + "@@F@@": RenderedToken.model_construct( + value="f", + kind="own_function", + vaddr=None, + data_type_id=None, + function_id=None, + imported_function_id=None, + ) + }, + placeholder_to_user_override={}, ) @@ -369,7 +379,7 @@ def stall_status(*a, **kw): def test_tokenised_phase_caches_and_dispatches_on_completed(service, sdk): sdk.get_ai_decompilation.return_value = _dd(code="ok") - sdk.get_ai_decompilation_tokenised.return_value = _tokd() + sdk.v3_get_ai_decompilation_tokens.return_value = _tokd() on_tokenised = MagicMock() service.start_ai_decomp_task( @@ -383,19 +393,42 @@ def test_tokenised_phase_caches_and_dispatches_on_completed(service, sdk): on_tokenised.assert_called_once() assert on_tokenised.call_args[0][0].success is True - assert service._tokenised_cache[42] is sdk.get_ai_decompilation_tokenised.return_value + assert service._tokenised_cache[42] is sdk.v3_get_ai_decompilation_tokens.return_value + + +def test_tokenised_phase_treats_an_empty_source_as_not_ready(service, sdk): + sdk.get_ai_decompilation.return_value = _dd(code="ok") + sdk.v3_get_ai_decompilation_tokens.return_value = GetTokensResponse.model_construct( + ai_decomp="", + analysis_id=1, + placeholder_to_rendered_token={}, + placeholder_to_user_override={}, + ) + + on_tokenised = MagicMock() + service.start_ai_decomp_task( + ea=4096, + on_decomp=MagicMock(), + on_summary=MagicMock(), + on_comments=MagicMock(), + on_tokenised=on_tokenised, + ) + _wait(service) + + assert on_tokenised.call_args[0][0].success is False + assert 42 not in service._tokenised_cache def test_tokenised_phase_skipped_when_no_callback(service, sdk): sdk.get_ai_decompilation.return_value = _dd(code="ok") _run(service) - sdk.get_ai_decompilation_tokenised.assert_not_called() + sdk.v3_get_ai_decompilation_tokens.assert_not_called() def test_apply_overrides_sends_body_refetches_and_caches(service, sdk): - sdk.upsert_ai_decompilation_overrides.return_value = MagicMock() + sdk.v3_upsert_ai_decompilation_overrides.return_value = MagicMock() sdk.get_ai_decompilation.return_value = _dd(code="renamed") - sdk.get_ai_decompilation_tokenised.return_value = _tokd() + sdk.v3_get_ai_decompilation_tokens.return_value = _tokd() on_decomp, on_tokenised = MagicMock(), MagicMock() service.apply_overrides( @@ -407,19 +440,19 @@ def test_apply_overrides_sends_body_refetches_and_caches(service, sdk): _wait_mock(on_decomp) _wait_mock(on_tokenised) - _, kwargs = sdk.upsert_ai_decompilation_overrides.call_args + _, kwargs = sdk.v3_upsert_ai_decompilation_overrides.call_args assert kwargs["function_id"] == 42 - assert kwargs["upsert_overrides_input_body"].overrides == {"@@V_v5@@": "buf"} + assert kwargs["upsert_overrides_input_body"].overrides == {"@@V_v5@@": Token(value="buf")} payload = on_decomp.call_args[0][0] assert payload.success is True assert payload.data.decompilation == "renamed" assert service._decomp_cache[42].decompilation == "renamed" - assert service._tokenised_cache[42] is sdk.get_ai_decompilation_tokenised.return_value + assert service._tokenised_cache[42] is sdk.v3_get_ai_decompilation_tokens.return_value def test_apply_overrides_api_error_surfaces(service, sdk): - sdk.upsert_ai_decompilation_overrides.side_effect = ApiException(status=500, reason="x") + sdk.v3_upsert_ai_decompilation_overrides.side_effect = ApiException(status=500, reason="x") on_decomp, on_tokenised = MagicMock(), MagicMock() service.apply_overrides( From 79c2d47df6e0f36e3a4981b6f4c88dcb8ddf3f79 Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:17:08 +0100 Subject: [PATCH 05/12] feat: stream ai-decomp --- .../app/components/tabs/ai_decomp_tab.py | 5 +- .../app/coordinators/ai_decomp_coordinator.py | 14 + .../app/coordinators/ai_decomp_render.py | 30 +++ .../services/ai_decomp/ai_decomp_service.py | 128 +++++++-- reai_toolkit/app/services/ai_decomp/stream.py | 216 +++++++++++++++ tests/unit/ai_decomp/test_render.py | 53 ++++ tests/unit/ai_decomp/test_service.py | 169 ++++++++++++ tests/unit/ai_decomp/test_stream.py | 253 ++++++++++++++++++ 8 files changed, 850 insertions(+), 18 deletions(-) create mode 100644 reai_toolkit/app/services/ai_decomp/stream.py create mode 100644 tests/unit/ai_decomp/test_stream.py diff --git a/reai_toolkit/app/components/tabs/ai_decomp_tab.py b/reai_toolkit/app/components/tabs/ai_decomp_tab.py index 5e2e63d..f1b8289 100644 --- a/reai_toolkit/app/components/tabs/ai_decomp_tab.py +++ b/reai_toolkit/app/components/tabs/ai_decomp_tab.py @@ -227,13 +227,16 @@ def _on_remove_comment_requested(self, line: int) -> None: # --- public API ------------------------------------------------ @execute_ui - def update_view_content(self, code: str) -> None: + def update_view_content(self, code: str, follow_tail: bool = False) -> None: if not self._editor: return self._editor.blockSignals(True) try: self._editor.setPlainText(code) + if follow_tail: + bar = self._editor.verticalScrollBar() + bar.setValue(bar.maximum()) finally: self._editor.blockSignals(False) diff --git a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py index 6868e3d..3c90e25 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py +++ b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py @@ -16,6 +16,7 @@ RenderModel, index_of_identifier, render_progress, + render_stream, render_view_with_map, resolve_token, ) @@ -23,6 +24,7 @@ from reai_toolkit.app.core.shared_schema import GenericApiReturn from reai_toolkit.app.factory import DialogFactory from reai_toolkit.app.services.ai_decomp.ai_decomp_service import AiDecompService +from reai_toolkit.app.services.ai_decomp.stream import StreamState from reai_toolkit.hooks.reactive import AiDecompFunctionViewHooks @@ -123,8 +125,10 @@ def _dispatch_task(self, ea: int) -> None: on_comments=lambda response: self._on_comments_complete(ea, response), on_tokenised=lambda response: self._on_tokenised_complete(ea, response), on_progress=lambda progress: self._on_progress(ea, progress), + on_stream=lambda state: self._on_stream(ea, state), ) + @execute_ui def _on_progress(self, ea: int, progress: WorkflowProgress) -> None: if ea != self._current_func_vaddr: return @@ -134,6 +138,16 @@ def _on_progress(self, ea: int, progress: WorkflowProgress) -> None: return self._decomp_view.update_view_content(render_progress(progress)) + @execute_ui + def _on_stream(self, ea: int, state: StreamState) -> None: + if ea != self._current_func_vaddr: + return + if self._current_decomp is not None: + return + if self._decomp_view is None: + return + self._decomp_view.update_view_content(render_stream(state), follow_tail=True) + def _on_decomp_complete( self, ea: int, response: GenericApiReturn[DecompilationData] ) -> None: diff --git a/reai_toolkit/app/coordinators/ai_decomp_render.py b/reai_toolkit/app/coordinators/ai_decomp_render.py index 8650483..f85addb 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_render.py +++ b/reai_toolkit/app/coordinators/ai_decomp_render.py @@ -4,6 +4,8 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Optional +from reai_toolkit.app.services.ai_decomp.stream import PROSE_TAIL + if TYPE_CHECKING: from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData @@ -13,6 +15,8 @@ from revengai.models.summary_data import SummaryData from revengai.models.workflow_progress import WorkflowProgress + from reai_toolkit.app.services.ai_decomp.stream import StreamState + _IDENT_RE = re.compile(r"[A-Za-z_]\w*") _RENAMED_ELSEWHERE_FIELDS = ("data_type_id", "function_id", "imported_function_id") @@ -114,6 +118,32 @@ def render_progress(progress: "WorkflowProgress") -> str: return "\n".join(lines) +def render_stream(state: "StreamState") -> str: + lines: list[str] = [] + + if state.failed: + lines.append("// RevEng.AI — AI decompilation failed") + if state.error: + lines.append(f"// {state.error}") + elif state.finished: + lines.append("// RevEng.AI — AI decompilation complete") + else: + attempt = f" (attempt {state.attempt})" if state.attempt > 1 else "" + stage = "naming identifiers" if state.decomp_finished else "decompiling" + lines.append(f"// RevEng.AI — {stage}…{attempt}") + + if state.prose and not state.source: + lines.append("//") + for text in state.prose[-PROSE_TAIL:]: + for part in text.split("\n"): + lines.append(f"// {part}") + + header = "\n".join(lines) + if not state.source: + return header + return f"{header}\n\n{state.source}" + + def _format_progress_message(message: "ProgressMessage") -> str: stamp = _format_progress_time(getattr(message, "timestamp", None)) prefix = f"{stamp} " if stamp else "" diff --git a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py index df961be..0929239 100644 --- a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py +++ b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py @@ -1,6 +1,8 @@ import threading +import time from typing import Any, Callable +import urllib3 from loguru import logger from revengai import ( ApiException, @@ -26,11 +28,17 @@ from reai_toolkit.app.core.shared_schema import GenericApiReturn from reai_toolkit.app.core.utils import parse_exception from reai_toolkit.app.interfaces.thread_service import IThreadService +from reai_toolkit.app.services.ai_decomp.stream import StreamState, iter_stream_events POLL_INTERVAL_SECONDS: float = 1.5 MAX_REQUEUE_ATTEMPTS = 2 +STREAM_CONNECT_TIMEOUT: float = 10.0 +STREAM_READ_TIMEOUT: float = 300.0 +STREAM_CHUNK_BYTES = 1024 +STREAM_DISPATCH_INTERVAL: float = 0.1 + class AiDecompService(IThreadService): def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration) -> None: @@ -41,6 +49,7 @@ def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration) self._tokenised_cache: dict[int, GetTokensResponse] = {} self._inflight: dict[int, threading.Event] = {} self._inflight_lock = threading.Lock() + self._active_streams: dict[int, urllib3.HTTPResponse] = {} def thread_in_progress(self) -> bool: return self.is_worker_running() @@ -55,6 +64,8 @@ def stop_worker(self) -> None: self._inflight.clear() for evt in events: evt.set() + for function_id in list(self._active_streams): + self._close_stream(function_id) self._decomp_cache.clear() self._summary_cache.clear() self._comments_cache.clear() @@ -78,6 +89,7 @@ def invalidate(self, function_id: int) -> None: evt = self._inflight.pop(function_id, None) if evt is not None: evt.set() + self._close_stream(function_id) def invalidate_ea(self, ea: int) -> None: function_id = self._get_function_id(start_ea=ea) @@ -92,6 +104,7 @@ def start_ai_decomp_task( on_comments: Callable[[GenericApiReturn[CommentsData]], None], on_tokenised: Callable[[GenericApiReturn[GetTokensResponse]], None] | None = None, on_progress: Callable[[WorkflowProgress], None] | None = None, + on_stream: Callable[[StreamState], None] | None = None, ) -> None: function_id = self._get_function_id(start_ea=ea) if function_id is None: @@ -114,6 +127,7 @@ def start_ai_decomp_task( on_comments, on_tokenised, on_progress, + on_stream, ), name=f"reai-aidecomp-{function_id}", daemon=True, @@ -360,12 +374,13 @@ def _run_task( on_comments: Callable[[GenericApiReturn[CommentsData]], None], on_tokenised: Callable[[GenericApiReturn[GetTokensResponse]], None] | None = None, on_progress: Callable[[WorkflowProgress], None] | None = None, + on_stream: Callable[[StreamState], None] | None = None, ) -> None: try: if stop_event.is_set(): return decomp = self._run_decomp_phase( - function_id, stop_event, on_decomp, on_progress + function_id, stop_event, on_decomp, on_progress, on_stream ) if decomp is None: return @@ -448,6 +463,7 @@ def _run_decomp_phase( stop_event: threading.Event, on_decomp: Callable[[GenericApiReturn[DecompilationData]], None], on_progress: Callable[[WorkflowProgress], None] | None = None, + on_stream: Callable[[StreamState], None] | None = None, ) -> DecompilationData | None: cached = self._decomp_cache.get(function_id) if cached is not None: @@ -490,25 +506,43 @@ def _run_decomp_phase( ) return None - polled_ok, poll_err = self._poll_workflow( - function_id=function_id, - stop_event=stop_event, - status_fn=self._fetch_decomp_status, - requeue_fn=self._queue_decompilation, - label="AI Decompilation", - on_progress=on_progress, + streamed = ( + self._stream_decomp(function_id, stop_event, on_stream) + if on_stream is not None and not stop_event.is_set() + else None ) - if not polled_ok: - if poll_err is not None: - self._safe_dispatch( - stop_event, - on_decomp, - GenericApiReturn[DecompilationData]( - success=False, error_message=poll_err - ), - ) + + if streamed is not None and streamed.failed: + self._safe_dispatch( + stop_event, + on_decomp, + GenericApiReturn[DecompilationData]( + success=False, + error_message=streamed.error or "AI decompilation failed.", + ), + ) return None + if streamed is None: + polled_ok, poll_err = self._poll_workflow( + function_id=function_id, + stop_event=stop_event, + status_fn=self._fetch_decomp_status, + requeue_fn=self._queue_decompilation, + label="AI Decompilation", + on_progress=on_progress, + ) + if not polled_ok: + if poll_err is not None: + self._safe_dispatch( + stop_event, + on_decomp, + GenericApiReturn[DecompilationData]( + success=False, error_message=poll_err + ), + ) + return None + final, final_err = self._fetch_decompilation(function_id) if final is None or not final.decompilation: self._safe_dispatch( @@ -529,6 +563,66 @@ def _run_decomp_phase( ) return final + def _close_stream(self, function_id: int) -> None: + resp = self._active_streams.pop(function_id, None) + if resp is None: + return + try: + resp.close() + except Exception: + pass + + def _stream_decomp( + self, + function_id: int, + stop_event: threading.Event, + on_stream: Callable[[StreamState], None], + ) -> StreamState | None: + state = StreamState() + last_dispatch = 0.0 + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + resp = FunctionsAIDecompilationApi( + api_client + ).stream_ai_decompilation_without_preload_content( + function_id=function_id, + _request_timeout=(STREAM_CONNECT_TIMEOUT, STREAM_READ_TIMEOUT), + ) + self._active_streams[function_id] = resp + try: + events = iter_stream_events( + resp.stream(STREAM_CHUNK_BYTES, decode_content=True), + stop=stop_event.is_set, + ) + for event in events: + state.apply(event) + due = time.monotonic() - last_dispatch >= STREAM_DISPATCH_INTERVAL + if state.is_terminal or due: + self._safe_dispatch(stop_event, on_stream, state.snapshot()) + last_dispatch = time.monotonic() + finally: + self._active_streams.pop(function_id, None) + try: + resp.release_conn() + except Exception: + pass + except (ApiException, urllib3.exceptions.HTTPError, OSError) as e: + logger.debug( + f"RevEng.AI: AI decompilation stream for function {function_id} ended: {e}" + ) + except Exception as e: + logger.debug( + f"RevEng.AI: AI decompilation stream for function {function_id} failed: {e}" + ) + + if state.is_terminal: + return state + logger.debug( + f"RevEng.AI: AI decompilation stream for function {function_id} " + "ended before a terminal event; falling back to polling" + ) + return None + def _queue_summary(self, function_id: int) -> tuple[bool, str | None]: try: with self.yield_api_client(sdk_config=self.sdk_config) as api_client: diff --git a/reai_toolkit/app/services/ai_decomp/stream.py b/reai_toolkit/app/services/ai_decomp/stream.py new file mode 100644 index 0000000..950aa51 --- /dev/null +++ b/reai_toolkit/app/services/ai_decomp/stream.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass, field, replace +from typing import Callable, Iterable, Iterator, Optional + +EVENT_ATTEMPT_FAILED = "attempt_failed" +EVENT_ATTEMPT_STARTED = "attempt_started" +EVENT_DECOMP_FAILED = "decomp_failed" +EVENT_DECOMP_FINISHED = "decomp_finished" +EVENT_NAMES_FINISHED = "names_finished" +EVENT_PROSE = "prose" +EVENT_RENAME_APPLIED = "rename_applied" +EVENT_SOURCE_DELTA = "source_delta" +EVENT_SOURCE_RESET = "source_reset" +EVENT_TYPES_SUGGESTED = "types_suggested" +EVENT_WARNING = "warning" + +TERMINAL_EVENTS: frozenset[str] = frozenset( + {EVENT_NAMES_FINISHED, EVENT_DECOMP_FAILED} +) + +PROSE_TAIL = 5 + + +@dataclass +class DecompEvent: + type: str + attempt: int = 0 + seq: int = 0 + content: str = "" + text: str = "" + error: str = "" + error_code: Optional[str] = None + old_name: str = "" + new_name: str = "" + warning_kind: str = "" + message: str = "" + applied: int = 0 + types: int = 0 + members: int = 0 + + +def _as_int(value) -> int: + return value if isinstance(value, int) and not isinstance(value, bool) else 0 + + +def _as_str(value) -> str: + return value if isinstance(value, str) else "" + + +def normalize_event(event_name: Optional[str], payload: dict) -> Optional[DecompEvent]: + event_type = _as_str(payload.get("type")) or _as_str(event_name) + if not event_type: + return None + error_code = payload.get("error_code") + return DecompEvent( + type=event_type, + attempt=_as_int(payload.get("attempt")), + seq=_as_int(payload.get("seq")), + content=_as_str(payload.get("content")), + text=_as_str(payload.get("text")), + error=_as_str(payload.get("error")), + error_code=error_code if isinstance(error_code, str) else None, + old_name=_as_str(payload.get("old_name")), + new_name=_as_str(payload.get("new_name")), + warning_kind=_as_str(payload.get("kind")), + message=_as_str(payload.get("message")), + applied=_as_int(payload.get("applied")), + types=_as_int(payload.get("types")), + members=_as_int(payload.get("members")), + ) + + +def _iter_lines( + chunks: Iterable[bytes], + stop: Optional[Callable[[], bool]] = None, +) -> Iterator[str]: + buf = b"" + for chunk in chunks: + if stop is not None and stop(): + return + if not chunk: + continue + buf += chunk + parts = buf.split(b"\n") + buf = parts.pop() + for raw in parts: + yield raw.rstrip(b"\r").decode("utf-8", "replace") + if buf: + yield buf.rstrip(b"\r").decode("utf-8", "replace") + + +def iter_sse_frames( + chunks: Iterable[bytes], + stop: Optional[Callable[[], bool]] = None, +) -> Iterator[tuple[Optional[str], str]]: + event_name: Optional[str] = None + data_lines: list[str] = [] + + for line in _iter_lines(chunks, stop=stop): + if not line: + if data_lines: + yield event_name, "\n".join(data_lines) + event_name = None + data_lines = [] + continue + if line.startswith(":"): + continue + name, _, value = line.partition(":") + if value.startswith(" "): + value = value[1:] + if name == "event": + event_name = value + elif name == "data": + data_lines.append(value) + + if data_lines: + yield event_name, "\n".join(data_lines) + + +def iter_stream_events( + chunks: Iterable[bytes], + stop: Optional[Callable[[], bool]] = None, +) -> Iterator[DecompEvent]: + for event_name, data in iter_sse_frames(chunks, stop=stop): + if data == "[DONE]": + return + try: + payload = json.loads(data) + except (ValueError, TypeError): + continue + if not isinstance(payload, dict): + continue + event = normalize_event(event_name, payload) + if event is None: + continue + yield event + if event.type in TERMINAL_EVENTS: + return + + +@dataclass +class StreamState: + attempt: int = 0 + source: str = "" + prose: list[str] = field(default_factory=list) + renames: list[tuple[str, str]] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + decomp_finished: bool = False + finished: bool = False + failed: bool = False + error: str = "" + error_code: Optional[str] = None + names_applied: int = 0 + types_suggested: int = 0 + members_suggested: int = 0 + + def reset_attempt(self, attempt: int) -> None: + self.attempt = attempt + self.source = "" + self.prose = [] + self.renames = [] + self.warnings = [] + self.decomp_finished = False + self.error = "" + self.error_code = None + + def apply(self, event: DecompEvent) -> None: + if event.attempt > self.attempt: + self.reset_attempt(event.attempt) + + if event.type == EVENT_SOURCE_DELTA: + self.source += event.content + elif event.type == EVENT_SOURCE_RESET: + self.source = "" + elif event.type == EVENT_PROSE: + if event.text: + self.prose.append(event.text) + elif event.type == EVENT_RENAME_APPLIED: + self.renames.append((event.old_name, event.new_name)) + elif event.type == EVENT_WARNING: + self.warnings.append(event.message or event.warning_kind) + elif event.type == EVENT_TYPES_SUGGESTED: + self.types_suggested = event.types + self.members_suggested = event.members + elif event.type == EVENT_DECOMP_FINISHED: + self.decomp_finished = True + elif event.type == EVENT_ATTEMPT_FAILED: + self.error = event.error + elif event.type == EVENT_DECOMP_FAILED: + self.error = event.error + self.error_code = event.error_code + self.failed = True + elif event.type == EVENT_NAMES_FINISHED: + self.names_applied = event.applied + self.finished = True + + @property + def is_terminal(self) -> bool: + return self.finished or self.failed + + def snapshot(self) -> "StreamState": + return replace( + self, + prose=list(self.prose), + renames=list(self.renames), + warnings=list(self.warnings), + ) + + +def reduce_stream(events: Iterable[DecompEvent]) -> StreamState: + state = StreamState() + for event in events: + state.apply(event) + return state diff --git a/tests/unit/ai_decomp/test_render.py b/tests/unit/ai_decomp/test_render.py index f5c2a33..baad964 100644 --- a/tests/unit/ai_decomp/test_render.py +++ b/tests/unit/ai_decomp/test_render.py @@ -15,10 +15,12 @@ index_of_identifier, names_token, render_progress, + render_stream, render_view, render_view_with_map, resolve_token, ) +from reai_toolkit.app.services.ai_decomp.stream import PROSE_TAIL, StreamState CODE = "int f(int a1) {\n int v5 = a1;\n return v5;\n}" @@ -235,6 +237,57 @@ def test_render_progress_lists_messages_as_comment_lines(): assert all(line.startswith("//") for line in text.split("\n")) +def _state(**kw): + return StreamState(**kw) + + +def test_render_stream_shows_the_source_as_it_arrives(): + text = render_stream(_state(attempt=1, source="int main(void) {")) + + assert text.startswith("// RevEng.AI — decompiling…") + assert text.endswith("int main(void) {") + + +def test_render_stream_shows_prose_only_until_source_starts(): + with_prose = _state(attempt=1, prose=["reading the bytes", "spotting a loop"]) + assert "// spotting a loop" in render_stream(with_prose) + + with_source = _state(attempt=1, prose=["reading the bytes"], source="int x;") + assert "reading the bytes" not in render_stream(with_source) + + +def test_render_stream_caps_the_prose_it_shows(): + state = _state(attempt=1, prose=[f"line {i}" for i in range(20)]) + + body = [line for line in render_stream(state).split("\n") if line.startswith("// line")] + assert len(body) == PROSE_TAIL + + +def test_render_stream_names_the_post_decompilation_naming_stage(): + text = render_stream(_state(attempt=1, source="int x;", decomp_finished=True)) + + assert "naming identifiers" in text + + +def test_render_stream_shows_the_attempt_only_after_a_retry(): + assert "attempt" not in render_stream(_state(attempt=1)) + assert "(attempt 2)" in render_stream(_state(attempt=2)) + + +def test_render_stream_reports_failure_with_its_error(): + text = render_stream(_state(failed=True, error="model unavailable")) + + assert "failed" in text + assert "// model unavailable" in text + + +def test_render_stream_of_a_finished_run_keeps_the_source(): + text = render_stream(_state(finished=True, source="int main(void) {}")) + + assert "complete" in text + assert text.endswith("int main(void) {}") + + def test_render_progress_without_steps_falls_back_to_status(): text = render_progress(_wp(step="", step_index=0, steps_total=0, status="PENDING")) assert "// PENDING" in text.split("\n") diff --git a/tests/unit/ai_decomp/test_service.py b/tests/unit/ai_decomp/test_service.py index 9d8212b..d3ea6ff 100644 --- a/tests/unit/ai_decomp/test_service.py +++ b/tests/unit/ai_decomp/test_service.py @@ -1,3 +1,5 @@ +import json +import pathlib import threading import time from unittest.mock import MagicMock @@ -581,3 +583,170 @@ def test_function_id_for_and_invalidate_ea(service): service._decomp_cache[42] = object() service.invalidate_ea(4096) assert 42 not in service._decomp_cache + + +def _sse(event_type, **payload): + body = {"type": event_type, "attempt": 1, "seq": payload.pop("seq", 0)} + body.update(payload) + return f"event: {event_type}\ndata: {json.dumps(body)}\n\n".encode() + + +def _stream_response(chunks): + resp = MagicMock() + resp.stream.return_value = iter(chunks) + return resp + + +def _start_streaming(service, on_stream, on_decomp=None): + service.start_ai_decomp_task( + ea=4096, + on_decomp=on_decomp or MagicMock(), + on_summary=MagicMock(), + on_comments=MagicMock(), + on_stream=on_stream, + ) + _wait(service) + + +def test_streaming_replaces_polling_and_reports_the_growing_source(service, sdk, monkeypatch): + monkeypatch.setattr(svc_mod, "STREAM_DISPATCH_INTERVAL", 0.0) + sdk.get_ai_decompilation.side_effect = [_dd(status=TaskStatus.RUNNING.value, code=None), _dd(code="final")] + sdk.stream_ai_decompilation_without_preload_content.return_value = _stream_response( + [ + _sse("source_delta", content="int main"), + _sse("source_delta", content="(void) {}"), + _sse("decomp_finished"), + _sse("names_finished", applied=2), + ] + ) + + on_stream, on_decomp = MagicMock(), MagicMock() + _start_streaming(service, on_stream, on_decomp) + + sdk.get_ai_decompilation_status.assert_not_called() + sources = [call[0][0].source for call in on_stream.call_args_list] + assert sources[-1] == "int main(void) {}" + assert on_stream.call_args[0][0].finished is True + assert on_decomp.call_args[0][0].data.decompilation == "final" + + +def test_each_dispatched_state_is_an_independent_snapshot(service, sdk, monkeypatch): + monkeypatch.setattr(svc_mod, "STREAM_DISPATCH_INTERVAL", 0.0) + sdk.get_ai_decompilation.side_effect = [_dd(status=TaskStatus.RUNNING.value, code=None), _dd(code="final")] + sdk.stream_ai_decompilation_without_preload_content.return_value = _stream_response( + [ + _sse("source_delta", content="a"), + _sse("source_delta", content="b"), + _sse("names_finished", applied=0), + ] + ) + + on_stream = MagicMock() + _start_streaming(service, on_stream) + + states = [call[0][0] for call in on_stream.call_args_list] + assert [s.source for s in states] == ["a", "ab", "ab"] + + +def test_a_stream_that_ends_without_a_terminal_event_falls_back_to_polling(service, sdk): + sdk.get_ai_decompilation.side_effect = [_dd(status=TaskStatus.RUNNING.value, code=None), _dd(code="polled")] + sdk.stream_ai_decompilation_without_preload_content.return_value = _stream_response( + [_sse("source_delta", content="partial")] + ) + sdk.get_ai_decompilation_status.return_value = _wp(TaskStatus.COMPLETED) + + on_decomp = MagicMock() + _start_streaming(service, MagicMock(), on_decomp) + + sdk.get_ai_decompilation_status.assert_called() + assert on_decomp.call_args[0][0].data.decompilation == "polled" + + +def test_a_stream_that_raises_falls_back_to_polling(service, sdk): + sdk.get_ai_decompilation.side_effect = [_dd(status=TaskStatus.RUNNING.value, code=None), _dd(code="polled")] + sdk.stream_ai_decompilation_without_preload_content.side_effect = ApiException( + status=502, reason="bad gateway" + ) + sdk.get_ai_decompilation_status.return_value = _wp(TaskStatus.COMPLETED) + + on_decomp = MagicMock() + _start_streaming(service, MagicMock(), on_decomp) + + sdk.get_ai_decompilation_status.assert_called() + assert on_decomp.call_args[0][0].data.decompilation == "polled" + + +def test_decomp_failed_surfaces_the_error_without_polling(service, sdk): + sdk.get_ai_decompilation.return_value = _dd(status=TaskStatus.RUNNING.value, code=None) + sdk.stream_ai_decompilation_without_preload_content.return_value = _stream_response( + [_sse("decomp_failed", error="model unavailable", error_code="MODEL_GONE")] + ) + + on_decomp = MagicMock() + _start_streaming(service, MagicMock(), on_decomp) + + sdk.get_ai_decompilation_status.assert_not_called() + payload = on_decomp.call_args[0][0] + assert payload.success is False + assert payload.error_message == "model unavailable" + + +def test_no_stream_callback_keeps_the_polling_path(service, sdk): + sdk.get_ai_decompilation.side_effect = [_dd(status=TaskStatus.RUNNING.value, code=None), _dd(code="polled")] + sdk.get_ai_decompilation_status.return_value = _wp(TaskStatus.COMPLETED) + + _run(service) + + sdk.stream_ai_decompilation_without_preload_content.assert_not_called() + sdk.get_ai_decompilation_status.assert_called() + + +def test_dispatches_are_throttled_between_terminal_events(service, sdk, monkeypatch): + monkeypatch.setattr(svc_mod, "STREAM_DISPATCH_INTERVAL", 60.0) + sdk.get_ai_decompilation.side_effect = [_dd(status=TaskStatus.RUNNING.value, code=None), _dd(code="final")] + sdk.stream_ai_decompilation_without_preload_content.return_value = _stream_response( + [_sse("source_delta", content=str(i)) for i in range(10)] + + [_sse("names_finished", applied=0)] + ) + + on_stream = MagicMock() + _start_streaming(service, on_stream) + + assert on_stream.call_count == 2 + assert on_stream.call_args[0][0].finished is True + + +def test_invalidate_closes_a_live_stream(service): + resp = MagicMock() + service._active_streams[42] = resp + + service.invalidate(42) + + resp.close.assert_called_once() + assert 42 not in service._active_streams + + +def test_streaming_never_dispatches_on_the_main_thread(service, sdk, monkeypatch): + monkeypatch.setattr(svc_mod, "STREAM_DISPATCH_INTERVAL", 0.0) + sdk.get_ai_decompilation.side_effect = [_dd(status=TaskStatus.RUNNING.value, code=None), _dd(code="final")] + sdk.stream_ai_decompilation_without_preload_content.return_value = _stream_response( + [_sse("source_delta", content="x"), _sse("names_finished", applied=0)] + ) + + threads: list = [] + _start_streaming(service, lambda state: threads.append(threading.current_thread())) + + assert threads + assert all(t is not threading.main_thread() for t in threads) + + +@pytest.mark.parametrize("callback", ["_on_stream", "_on_progress"]) +def test_view_writing_callbacks_stay_marshalled_onto_the_ui_thread(callback): + source = pathlib.Path( + "reai_toolkit/app/coordinators/ai_decomp_coordinator.py" + ).read_text() + + assert f" @execute_ui\n def {callback}(" in source, ( + f"{callback} writes to the Qt view from a stream/poll worker thread, " + "so it must keep @execute_ui" + ) diff --git a/tests/unit/ai_decomp/test_stream.py b/tests/unit/ai_decomp/test_stream.py new file mode 100644 index 0000000..226128f --- /dev/null +++ b/tests/unit/ai_decomp/test_stream.py @@ -0,0 +1,253 @@ +import json + +import pytest + +from reai_toolkit.app.services.ai_decomp.stream import ( + EVENT_ATTEMPT_FAILED, + EVENT_ATTEMPT_STARTED, + EVENT_DECOMP_FAILED, + EVENT_DECOMP_FINISHED, + EVENT_NAMES_FINISHED, + EVENT_PROSE, + EVENT_RENAME_APPLIED, + EVENT_SOURCE_DELTA, + EVENT_SOURCE_RESET, + EVENT_TYPES_SUGGESTED, + EVENT_WARNING, + DecompEvent, + StreamState, + iter_sse_frames, + iter_stream_events, + normalize_event, + reduce_stream, +) + + +def _frame(event_type, **payload): + body = {"type": event_type, "attempt": 1, "seq": payload.pop("seq", 0)} + body.update(payload) + return f"event: {event_type}\ndata: {json.dumps(body)}\n\n".encode() + + +def _ev(event_type, **kw): + kw.setdefault("attempt", 1) + return DecompEvent(type=event_type, **kw) + + +def test_frames_split_on_blank_lines_and_keep_the_event_name(): + chunks = [_frame(EVENT_SOURCE_DELTA, content="int "), _frame(EVENT_DECOMP_FINISHED)] + + frames = list(iter_sse_frames(chunks)) + + assert [name for name, _ in frames] == [EVENT_SOURCE_DELTA, EVENT_DECOMP_FINISHED] + + +def test_frames_survive_being_split_across_chunk_boundaries(): + raw = _frame(EVENT_SOURCE_DELTA, content="int main(void)") + chunks = [raw[:7], raw[7:20], raw[20:]] + + events = list(iter_stream_events(chunks)) + + assert [e.content for e in events] == ["int main(void)"] + + +def test_multiline_data_lines_are_joined_with_newlines(): + payload = json.dumps({"type": EVENT_PROSE, "attempt": 1, "seq": 1, "text": "hi"}) + half = len(payload) // 2 + raw = f"data: {payload[:half]}\ndata: {payload[half:]}\n\n".encode() + + frames = list(iter_sse_frames([raw])) + + assert frames == [(None, f"{payload[:half]}\n{payload[half:]}")] + + +def test_comment_lines_and_unparseable_frames_are_skipped(): + chunks = [ + b": keep-alive\n\n", + b"data: not json\n\n", + b"data: [1, 2, 3]\n\n", + _frame(EVENT_SOURCE_DELTA, content="ok"), + ] + + events = list(iter_stream_events(chunks)) + + assert [e.content for e in events] == ["ok"] + + +def test_done_sentinel_ends_the_stream(): + chunks = [_frame(EVENT_SOURCE_DELTA, content="a"), b"data: [DONE]\n\n", _frame(EVENT_PROSE, text="never")] + + events = list(iter_stream_events(chunks)) + + assert [e.type for e in events] == [EVENT_SOURCE_DELTA] + + +def test_a_trailing_frame_without_its_blank_line_is_still_yielded(): + raw = _frame(EVENT_SOURCE_DELTA, content="tail")[:-2] + + events = list(iter_stream_events([raw])) + + assert [e.content for e in events] == ["tail"] + + +@pytest.mark.parametrize("terminal", [EVENT_NAMES_FINISHED, EVENT_DECOMP_FAILED]) +def test_iteration_stops_after_a_terminal_event(terminal): + chunks = [_frame(terminal, error="boom"), _frame(EVENT_PROSE, text="after the end")] + + events = list(iter_stream_events(chunks)) + + assert [e.type for e in events] == [terminal] + + +def test_decomp_finished_is_not_terminal(): + chunks = [_frame(EVENT_DECOMP_FINISHED), _frame(EVENT_NAMES_FINISHED, applied=3)] + + events = list(iter_stream_events(chunks)) + + assert [e.type for e in events] == [EVENT_DECOMP_FINISHED, EVENT_NAMES_FINISHED] + + +def test_stop_callback_halts_iteration(): + calls = {"n": 0} + + def chunks(): + for _ in range(5): + yield _frame(EVENT_SOURCE_DELTA, content="x") + + def stop(): + calls["n"] += 1 + return calls["n"] > 2 + + assert len(list(iter_stream_events(chunks(), stop=stop))) == 2 + + +def test_type_falls_back_to_the_sse_event_line(): + event = normalize_event(EVENT_SOURCE_DELTA, {"content": "x", "attempt": 1}) + + assert event.type == EVENT_SOURCE_DELTA + assert event.content == "x" + + +def test_payload_type_wins_over_the_event_line(): + event = normalize_event("stale", {"type": EVENT_PROSE, "text": "x"}) + + assert event.type == EVENT_PROSE + + +def test_an_event_with_no_type_anywhere_is_dropped(): + assert normalize_event(None, {"content": "x"}) is None + + +def test_unknown_event_types_survive_normalisation(): + event = normalize_event(None, {"type": "invented_later", "attempt": 2}) + + assert event.type == "invented_later" + assert event.attempt == 2 + + +def test_wrongly_typed_fields_do_not_crash_normalisation(): + event = normalize_event(None, {"type": EVENT_SOURCE_DELTA, "content": 5, "attempt": "x"}) + + assert event.content == "" + assert event.attempt == 0 + + +def test_source_deltas_accumulate(): + state = reduce_stream( + [ + _ev(EVENT_SOURCE_DELTA, content="int main"), + _ev(EVENT_SOURCE_DELTA, content="(void) {}"), + ] + ) + + assert state.source == "int main(void) {}" + + +def test_source_reset_clears_the_buffer_without_touching_the_attempt(): + state = reduce_stream( + [ + _ev(EVENT_SOURCE_DELTA, content="discard me"), + _ev(EVENT_SOURCE_RESET), + _ev(EVENT_SOURCE_DELTA, content="keep me"), + ] + ) + + assert state.source == "keep me" + assert state.attempt == 1 + + +def test_a_new_attempt_discards_everything_from_the_previous_one(): + state = reduce_stream( + [ + _ev(EVENT_SOURCE_DELTA, content="first try"), + _ev(EVENT_PROSE, text="thinking"), + _ev(EVENT_ATTEMPT_FAILED, error="flaky"), + _ev(EVENT_ATTEMPT_STARTED, attempt=2), + _ev(EVENT_SOURCE_DELTA, attempt=2, content="second try"), + ] + ) + + assert state.source == "second try" + assert state.prose == [] + assert state.error == "" + assert state.attempt == 2 + + +def test_attempt_failed_is_recorded_but_not_terminal(): + state = reduce_stream([_ev(EVENT_ATTEMPT_FAILED, error="flaky")]) + + assert state.error == "flaky" + assert state.is_terminal is False + + +def test_decomp_finished_marks_the_naming_stage_not_completion(): + state = reduce_stream([_ev(EVENT_DECOMP_FINISHED)]) + + assert state.decomp_finished is True + assert state.finished is False + assert state.is_terminal is False + + +def test_names_finished_completes_the_run(): + state = reduce_stream([_ev(EVENT_DECOMP_FINISHED), _ev(EVENT_NAMES_FINISHED, applied=4)]) + + assert state.finished is True + assert state.is_terminal is True + assert state.names_applied == 4 + + +def test_decomp_failed_records_the_error_and_terminates(): + state = reduce_stream([_ev(EVENT_DECOMP_FAILED, error="no model", error_code="MODEL_GONE")]) + + assert (state.failed, state.error, state.error_code) == (True, "no model", "MODEL_GONE") + assert state.is_terminal is True + + +def test_renames_warnings_and_type_suggestions_are_collected(): + state = reduce_stream( + [ + _ev(EVENT_RENAME_APPLIED, old_name="v1", new_name="count"), + _ev(EVENT_WARNING, message="truncated output", warning_kind="TRUNCATED"), + _ev(EVENT_TYPES_SUGGESTED, types=2, members=7), + ] + ) + + assert state.renames == [("v1", "count")] + assert state.warnings == ["truncated output"] + assert (state.types_suggested, state.members_suggested) == (2, 7) + + +def test_a_warning_with_no_message_falls_back_to_its_kind(): + state = reduce_stream([_ev(EVENT_WARNING, warning_kind="TRUNCATED")]) + + assert state.warnings == ["TRUNCATED"] + + +def test_empty_prose_is_not_recorded(): + state = reduce_stream([_ev(EVENT_PROSE, text="")]) + + assert state.prose == [] + + +def test_a_fresh_state_is_not_terminal(): + assert StreamState().is_terminal is False From e4c282969344b65427974a2f87765cd86f344da9 Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:40:10 +0100 Subject: [PATCH 06/12] feat: predicted name and rename decline reasons --- .../app/components/tabs/ai_decomp_tab.py | 47 +++++++++- .../app/coordinators/ai_decomp_coordinator.py | 52 ++++++----- .../app/coordinators/ai_decomp_render.py | 69 +++++++++++++- tests/ui/ai_decomp_ui_runner.py | 93 ++++++++++++++++++- tests/ui/test_gui_ai_decomp.py | 8 ++ tests/unit/ai_decomp/test_render.py | 82 ++++++++++++++++ tests/unit/ai_decomp/test_service.py | 36 +++++++ 7 files changed, 357 insertions(+), 30 deletions(-) diff --git a/reai_toolkit/app/components/tabs/ai_decomp_tab.py b/reai_toolkit/app/components/tabs/ai_decomp_tab.py index f1b8289..80bd836 100644 --- a/reai_toolkit/app/components/tabs/ai_decomp_tab.py +++ b/reai_toolkit/app/components/tabs/ai_decomp_tab.py @@ -1,3 +1,4 @@ +import re from pathlib import Path from typing import Any, Callable, Optional from loguru import logger @@ -12,6 +13,12 @@ getattr(QtGui.QTextCursor, "SelectionType", QtGui.QTextCursor), "WordUnderCursor" ) +_IDENTIFIER_RE = re.compile(r"[A-Za-z_]\w*\Z") + + +def _is_identifier(word: str) -> bool: + return bool(word) and _IDENTIFIER_RE.match(word) is not None + def _thumb_icon(up: bool) -> Optional[QtGui.QIcon]: stem = "thumb_up" if up else "thumb_down" @@ -38,7 +45,7 @@ def mouseDoubleClickEvent(self, event) -> None: super().mouseDoubleClickEvent(event) cursor = self.textCursor() word = cursor.selectedText() - if word: + if _is_identifier(word): self.renameRequested.emit(cursor.blockNumber(), word) def contextMenuEvent(self, event) -> None: @@ -49,7 +56,7 @@ def contextMenuEvent(self, event) -> None: menu = self.createStandardContextMenu() menu.addSeparator() - act_rename = menu.addAction(f"Rename '{word}'…") if word else None + act_rename = menu.addAction(f"Rename '{word}'…") if _is_identifier(word) else None act_comment = menu.addAction("Add / edit comment…") act_remove = menu.addAction("Remove comment") @@ -85,11 +92,15 @@ def __init__(self, on_closed: Optional[Callable[[], None]] = None) -> None: self.on_remove_comment: Callable[[int], None] | None = None self.on_rate_up: Callable[[], None] | None = None self.on_rate_down: Callable[[], None] | None = None + self.on_use_predicted_name: Callable[[str], None] | None = None self._parent_window: QtWidgets.QWidget | None = None self._editor: _DecompEditor | None = None self._refresh_btn: QtWidgets.QPushButton | None = None self._rate_up_btn: QtWidgets.QPushButton | None = None self._rate_down_btn: QtWidgets.QPushButton | None = None + self._predicted_label: QtWidgets.QLabel | None = None + self._predicted_btn: QtWidgets.QPushButton | None = None + self._predicted_name: str | None = None self._highlighter: CppHighlighter | None = None def Create(self, title: Any) -> Any: @@ -122,6 +133,19 @@ def OnCreate(self, form) -> None: header = QtWidgets.QHBoxLayout() title = QtWidgets.QLabel("RevEng.AI — AI Decomp", self._parent_window) header.addWidget(title) + + self._predicted_label = QtWidgets.QLabel("", self._parent_window) + self._predicted_label.setVisible(False) + header.addWidget(self._predicted_label) + + self._predicted_btn = QtWidgets.QPushButton( + "Use Predicted Name", self._parent_window + ) + self._predicted_btn.setToolTip("Rename this function to the predicted name") + self._predicted_btn.setVisible(False) + self._predicted_btn.clicked.connect(self._on_use_predicted_name_clicked) + header.addWidget(self._predicted_btn) + header.addStretch(1) up_icon = _thumb_icon(up=True) @@ -190,6 +214,9 @@ def OnClose(self, form) -> None: self._refresh_btn = None self._rate_up_btn = None self._rate_down_btn = None + self._predicted_label = None + self._predicted_btn = None + self._predicted_name = None self._parent_window = None def _on_refresh_clicked(self) -> None: @@ -206,6 +233,10 @@ def _on_rate_down_clicked(self) -> None: if self.on_rate_down: self.on_rate_down() + def _on_use_predicted_name_clicked(self) -> None: + if self._predicted_name and self.on_use_predicted_name: + self.on_use_predicted_name(self._predicted_name) + @execute_ui def set_rating(self, rating: Optional[str]) -> None: if self._rate_up_btn: @@ -213,6 +244,18 @@ def set_rating(self, rating: Optional[str]) -> None: if self._rate_down_btn: self._rate_down_btn.setChecked(rating == "down") + @execute_ui + def set_predicted_name(self, name: Optional[str]) -> None: + self._predicted_name = name or None + has_name = self._predicted_name is not None + if self._predicted_label: + self._predicted_label.setText( + f"Predicted name: {self._predicted_name}" if has_name else "" + ) + self._predicted_label.setVisible(has_name) + if self._predicted_btn: + self._predicted_btn.setVisible(has_name) + def _on_rename_requested(self, line: int, word: str) -> None: if self.on_rename: self.on_rename(line, word) diff --git a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py index 3c90e25..45dfbd4 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py +++ b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py @@ -13,12 +13,12 @@ from reai_toolkit.app.app import App from reai_toolkit.app.components.tabs.ai_decomp_tab import AIDecompView from reai_toolkit.app.coordinators.ai_decomp_render import ( + RENAME_NOT_READY, RenderModel, - index_of_identifier, render_progress, render_stream, render_view_with_map, - resolve_token, + resolve_rename_target, ) from reai_toolkit.app.coordinators.base_coordinator import BaseCoordinator from reai_toolkit.app.core.shared_schema import GenericApiReturn @@ -75,6 +75,7 @@ def run_dialog(self) -> None: self._decomp_view.on_remove_comment = self.request_remove_comment self._decomp_view.on_rate_up = self.rate_up self._decomp_view.on_rate_down = self.rate_down + self._decomp_view.on_use_predicted_name = self.apply_predicted_name self._decomp_view.Create(self._decomp_view.TITLE) def start_decompilation(self, ea: int) -> None: @@ -90,6 +91,7 @@ def start_decompilation(self, ea: int) -> None: cached = self.ai_decomp_service.peek_decomp(ea) if self._decomp_view is not None: self._decomp_view.set_rating(None) + self._decomp_view.set_predicted_name(None) if cached is not None and cached.decompilation: self._current_decomp = cached self._rerender() @@ -190,6 +192,8 @@ def _on_summary_complete( if response.data is None: return self._current_summary = response.data + if self._decomp_view is not None: + self._decomp_view.set_predicted_name(response.data.predicted_function_name) self._rerender() def _on_comments_complete( @@ -232,7 +236,7 @@ def _rate(self, rating: AiDecompilationRating) -> None: if ea is None: return if self._current_decomp is None: - self.show_info_dialog(message="No AI decompilation to rate yet.") + self.show_info_dialog(msg="No AI decompilation to rate yet.") if self._decomp_view is not None: self._decomp_view.set_rating(None) return @@ -253,39 +257,45 @@ def _on_rating_result( if self._decomp_view is not None: self._decomp_view.set_rating(None) - def request_rename(self, display_line: int, word: str) -> None: + def apply_predicted_name(self, name: str) -> None: ea = self._current_func_vaddr - if ea is None or self._baseline is None or self._current_tokenised is None: + if ea is None or not name: return - if not (0 <= display_line < len(self._baseline.display_is_code)): + + if self.ai_decomp_service.update_function_name(ea, name): + self.ai_decomp_service.tag_function_as_renamed(name) + self.refresh_disassembly_view() return - if not self._baseline.display_is_code[display_line]: + + final = self.ai_decomp_service.apply_deduped_name(ea, name) + if final is None: + self.show_info_dialog(msg=f"Could not rename this function to '{name}'.") return + self.ai_decomp_service.tag_function_as_renamed(final) + self.refresh_disassembly_view() - source_line = self._baseline.display_source[display_line] - if source_line is None: + def request_rename(self, display_line: int, word: str) -> None: + ea = self._current_func_vaddr + if ea is None: return - source_index = source_line - 1 - code_line = self._baseline.code_lines[source_index] - ident_index = index_of_identifier(code_line, word) - if ident_index < 0: + if self._baseline is None or self._current_tokenised is None: + self.show_info_dialog(msg=RENAME_NOT_READY) return - resolved = resolve_token(self._current_tokenised, source_index, ident_index, word) - if resolved is None: - self.show_info_dialog( - message=f"'{word}' is not a renameable variable or type." - ) + target = resolve_rename_target( + self._baseline, self._current_tokenised, display_line, word + ) + if target.placeholder is None: + self.show_info_dialog(msg=target.reason) return - token, _category = resolved new_name = ida_kernwin.ask_str(word, 0, f"Rename '{word}'") if not new_name or new_name == word: return self.ai_decomp_service.apply_overrides( ea=ea, - overrides={token: new_name}, + overrides={target.placeholder: new_name}, on_decomp=lambda response: self._on_decomp_complete(ea, response), on_tokenised=lambda response: self._on_tokenised_complete(ea, response), ) @@ -327,7 +337,7 @@ def request_remove_comment(self, display_line: int) -> None: if source_line is None: return if source_line not in self._baseline.comment_by_source: - self.show_info_dialog(message="No comment on this line.") + self.show_info_dialog(msg="No comment on this line.") return self.ai_decomp_service.remove_comment( ea=ea, diff --git a/reai_toolkit/app/coordinators/ai_decomp_render.py b/reai_toolkit/app/coordinators/ai_decomp_render.py index f85addb..5d4105a 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_render.py +++ b/reai_toolkit/app/coordinators/ai_decomp_render.py @@ -19,7 +19,27 @@ _IDENT_RE = re.compile(r"[A-Za-z_]\w*") -_RENAMED_ELSEWHERE_FIELDS = ("data_type_id", "function_id", "imported_function_id") + +RENAME_NOT_READY = "The AI decompilation is still loading — try again once it finishes." +RENAME_NOT_DECOMP_LINE = "This line is not part of the AI decompilation." +RENAME_NOT_CODE_LINE = ( + "Only identifiers in the decompiled code can be renamed, not comments." +) +RENAME_NOT_ON_LINE = "'{word}' is not an identifier on this line." +RENAME_IS_DATA_TYPE = ( + "'{word}' is a data type — rename the type itself, not the AI decompilation." +) +RENAME_IS_FUNCTION = ( + "'{word}' is a function — rename it in the disassembly and it will sync." +) +RENAME_IS_IMPORTED_FUNCTION = "'{word}' is an imported function and cannot be renamed." +RENAME_UNRESOLVED = "'{word}' is not a renameable variable or type." + +_RENAMED_ELSEWHERE: tuple[tuple[str, str], ...] = ( + ("data_type_id", RENAME_IS_DATA_TYPE), + ("function_id", RENAME_IS_FUNCTION), + ("imported_function_id", RENAME_IS_IMPORTED_FUNCTION), +) @dataclass @@ -31,6 +51,13 @@ class RenderModel: display_is_code: list[bool] +@dataclass +class RenameTarget: + placeholder: Optional[str] = None + kind: Optional[str] = None + reason: Optional[str] = None + + def render_view( decomp: "DecompilationData", summary: "Optional[SummaryData]", @@ -240,8 +267,15 @@ def find_token( return None +def renamed_elsewhere_reason(token: "RenderedToken", word: str) -> Optional[str]: + for name, reason in _RENAMED_ELSEWHERE: + if getattr(token, name, None) is not None: + return reason.format(word=word) + return None + + def is_renameable(token: "RenderedToken") -> bool: - return all(getattr(token, name, None) is None for name in _RENAMED_ELSEWHERE_FIELDS) + return all(getattr(token, name, None) is None for name, _ in _RENAMED_ELSEWHERE) def resolve_token( @@ -257,3 +291,34 @@ def resolve_token( if not is_renameable(token): return None return placeholder, token.kind + + +def resolve_rename_target( + model: RenderModel, + tokens: "GetTokensResponse", + display_line: int, + word: str, +) -> RenameTarget: + if not (0 <= display_line < len(model.display_is_code)): + return RenameTarget(reason=RENAME_NOT_DECOMP_LINE) + if not model.display_is_code[display_line]: + return RenameTarget(reason=RENAME_NOT_CODE_LINE) + + source_line = model.display_source[display_line] + if source_line is None: + return RenameTarget(reason=RENAME_NOT_DECOMP_LINE) + + source_index = source_line - 1 + ident_index = index_of_identifier(model.code_lines[source_index], word) + if ident_index < 0: + return RenameTarget(reason=RENAME_NOT_ON_LINE.format(word=word)) + + found = find_token(tokens, source_index, ident_index, word) + if found is None: + return RenameTarget(reason=RENAME_UNRESOLVED.format(word=word)) + + placeholder, token = found + reason = renamed_elsewhere_reason(token, word) + if reason is not None: + return RenameTarget(reason=reason) + return RenameTarget(placeholder=placeholder, kind=token.kind) diff --git a/tests/ui/ai_decomp_ui_runner.py b/tests/ui/ai_decomp_ui_runner.py index 4a92feb..0116698 100644 --- a/tests/ui/ai_decomp_ui_runner.py +++ b/tests/ui/ai_decomp_ui_runner.py @@ -40,8 +40,15 @@ def _run(report: dict) -> None: from revengai.models.inline_comment import InlineComment from revengai.models.rendered_token import RenderedToken - from reai_toolkit.app.components.tabs.ai_decomp_tab import AIDecompView + from revengai.models.summary_data import SummaryData + + from reai_toolkit.app.components.tabs.ai_decomp_tab import AIDecompView, _is_identifier from reai_toolkit.app.coordinators.ai_decomp_coordinator import AiDecompCoordinator + from reai_toolkit.app.coordinators.ai_decomp_render import ( + RENAME_IS_DATA_TYPE, + RENAME_NOT_CODE_LINE, + RENAME_UNRESOLVED, + ) from reai_toolkit.app.core.qt_compat import QtWidgets from reai_toolkit.app.core.shared_schema import GenericApiReturn @@ -52,14 +59,14 @@ def _run(report: dict) -> None: ea = next(iter(idautils.Functions()), 0x1000) decomp = DecompilationData.model_construct(status="COMPLETED", decompilation=CODE) - def _rt(value, kind): + def _rt(value, kind, **ids): return RenderedToken.model_construct( value=value, kind=kind, vaddr=None, - data_type_id=None, - function_id=None, - imported_function_id=None, + data_type_id=ids.get("data_type_id"), + function_id=ids.get("function_id"), + imported_function_id=ids.get("imported_function_id"), ) tokenised = GetTokensResponse.model_construct( @@ -72,6 +79,16 @@ def _rt(value, kind): placeholder_to_user_override={}, ) + typed = GetTokensResponse.model_construct( + ai_decomp="@@T@@ *x;", + analysis_id=1, + placeholder_to_rendered_token={"@@T@@": _rt("Foo", "type", data_type_id=7)}, + placeholder_to_user_override={}, + ) + typed_decomp = DecompilationData.model_construct( + status="COMPLETED", decompilation="Foo *x;" + ) + service = MagicMock() service.peek_decomp.return_value = None factory = SimpleNamespace( @@ -99,6 +116,12 @@ def pump() -> None: return report["view_created"] = True report["editor_read_only"] = view._editor.isReadOnly() is True + report["only_identifiers_offer_rename"] = ( + _is_identifier("v5") + and not _is_identifier(" ") + and not _is_identifier(";") + and not _is_identifier("") + ) def seed_plain() -> None: service.reset_mock() @@ -152,6 +175,58 @@ def code_line_row(needle: str) -> int: view._editor.renameRequested.emit(0, "int") pump() report["rename_non_token_info"] = not service.apply_overrides.called and len(infos) >= 1 + report["rename_non_token_reason"] = infos[:1] == [ + {"msg": RENAME_UNRESOLVED.format(word="int")} + ] + + seed_with_comment() + infos.clear() + view._editor.renameRequested.emit(code_line_row("// hola"), "hola") + pump() + report["rename_comment_line_reason"] = infos[:1] == [{"msg": RENAME_NOT_CODE_LINE}] + + seed_plain() + coord._current_decomp = typed_decomp + coord._current_tokenised = typed + coord._rerender() + pump() + infos.clear() + view._editor.renameRequested.emit(code_line_row("Foo *x;"), "Foo") + pump() + report["rename_data_type_reason"] = infos[:1] == [ + {"msg": RENAME_IS_DATA_TYPE.format(word="Foo")} + ] + + seed_plain() + report["predicted_hidden_without_a_prediction"] = view._predicted_btn.isHidden() + coord._on_summary_complete( + ea, + GenericApiReturn( + success=True, + data=SummaryData.model_construct( + ai_summary=None, + summary=None, + predicted_function_name="do_thing", + task_status="COMPLETED", + ), + ), + ) + pump() + report["predicted_shown_with_a_prediction"] = ( + not view._predicted_btn.isHidden() + and "do_thing" in view._predicted_label.text() + ) + + view._predicted_btn.click() + pump() + report["predicted_button_renames"] = ( + service.update_function_name.call_args is not None + and service.update_function_name.call_args.args == (ea, "do_thing") + ) + + view.set_predicted_name(None) + pump() + report["predicted_hidden_when_cleared"] = view._predicted_btn.isHidden() seed_plain() answers["text"] = "hello" @@ -198,10 +273,18 @@ def main() -> None: "errors": [], "view_created": False, "editor_read_only": False, + "only_identifiers_offer_rename": False, "render_shows_code": False, "rename_double_click_overrides": False, "rename_overrides_correct": False, "rename_non_token_info": False, + "rename_non_token_reason": False, + "rename_comment_line_reason": False, + "rename_data_type_reason": False, + "predicted_hidden_without_a_prediction": False, + "predicted_shown_with_a_prediction": False, + "predicted_button_renames": False, + "predicted_hidden_when_cleared": False, "comment_add_sets": False, "comment_add_args_correct": False, "comment_edit_empty_removes": False, diff --git a/tests/ui/test_gui_ai_decomp.py b/tests/ui/test_gui_ai_decomp.py index c608495..8f2d592 100644 --- a/tests/ui/test_gui_ai_decomp.py +++ b/tests/ui/test_gui_ai_decomp.py @@ -15,10 +15,18 @@ CHECKS = [ "view_created", "editor_read_only", + "only_identifiers_offer_rename", "render_shows_code", "rename_double_click_overrides", "rename_overrides_correct", "rename_non_token_info", + "rename_non_token_reason", + "rename_comment_line_reason", + "rename_data_type_reason", + "predicted_hidden_without_a_prediction", + "predicted_shown_with_a_prediction", + "predicted_button_renames", + "predicted_hidden_when_cleared", "comment_add_sets", "comment_add_args_correct", "comment_edit_empty_removes", diff --git a/tests/unit/ai_decomp/test_render.py b/tests/unit/ai_decomp/test_render.py index baad964..4fae7af 100644 --- a/tests/unit/ai_decomp/test_render.py +++ b/tests/unit/ai_decomp/test_render.py @@ -10,6 +10,13 @@ from revengai.models.workflow_progress import WorkflowProgress from reai_toolkit.app.coordinators.ai_decomp_render import ( + RENAME_IS_DATA_TYPE, + RENAME_IS_FUNCTION, + RENAME_IS_IMPORTED_FUNCTION, + RENAME_NOT_CODE_LINE, + RENAME_NOT_DECOMP_LINE, + RENAME_NOT_ON_LINE, + RENAME_UNRESOLVED, effective_values, find_token, index_of_identifier, @@ -18,6 +25,7 @@ render_stream, render_view, render_view_with_map, + resolve_rename_target, resolve_token, ) from reai_toolkit.app.services.ai_decomp.stream import PROSE_TAIL, StreamState @@ -207,6 +215,80 @@ def test_find_token_returns_tokens_renamed_elsewhere_even_though_resolve_decline assert token.data_type_id == 12 +def _model(code=CODE, summary=None, comments=None): + _, model = render_view_with_map(_dd(code), summary, comments) + return model + + +def test_rename_target_resolves_a_variable_on_a_code_line(): + target = resolve_rename_target(_model(), _tokens(rendered=_VARS), 1, "v5") + + assert target.placeholder == "@@V_v5@@" + assert target.kind == "local" + assert target.reason is None + + +def test_rename_target_declines_a_line_outside_the_view(): + for display_line in (-1, 99): + target = resolve_rename_target(_model(), _tokens(rendered=_VARS), display_line, "v5") + assert target.placeholder is None + assert target.reason == RENAME_NOT_DECOMP_LINE + + +def test_rename_target_declines_an_inline_comment_line(): + model = _model(comments=_comments([(2, "note")])) + target = resolve_rename_target(model, _tokens(rendered=_VARS), 1, "note") + + assert target.reason == RENAME_NOT_CODE_LINE + + +def test_rename_target_declines_a_summary_line(): + model = _model(summary=_summary("Adds one.")) + target = resolve_rename_target(model, _tokens(rendered=_VARS), 0, "Adds") + + assert target.reason == RENAME_NOT_CODE_LINE + + +def test_rename_target_declines_a_word_that_is_not_on_that_source_line(): + target = resolve_rename_target(_model(), _tokens(rendered=_VARS), 2, "a1") + + assert target.reason == RENAME_NOT_ON_LINE.format(word="a1") + + +def test_rename_target_declines_an_identifier_with_no_token(): + target = resolve_rename_target(_model(), _tokens(rendered=_VARS), 1, "int") + + assert target.reason == RENAME_UNRESOLVED.format(word="int") + + +@pytest.mark.parametrize( + "field,reason", + [ + ("data_type_id", RENAME_IS_DATA_TYPE), + ("function_id", RENAME_IS_FUNCTION), + ("imported_function_id", RENAME_IS_IMPORTED_FUNCTION), + ], +) +def test_rename_target_says_where_a_token_renamed_elsewhere_belongs(field, reason): + tokens = _tokens(tok="@@X@@ *x;", rendered={"@@X@@": _rt("Foo", kind="type", **{field: 12})}) + target = resolve_rename_target(_model(code="Foo *x;"), tokens, 0, "Foo") + + assert target.placeholder is None + assert target.reason == reason.format(word="Foo") + + +def test_every_declined_rename_carries_a_reason(): + model = _model(summary=_summary("S."), comments=_comments([(2, "note")])) + tokens = _tokens(rendered=_VARS) + + declines = [ + resolve_rename_target(model, tokens, line, word) + for line, word in [(-1, "v5"), (0, "S"), (99, "v5"), (5, "int"), (5, "nope")] + ] + + assert all(t.placeholder is None and t.reason for t in declines) + + def _pm(text, level="INFO", step="DECOMPILING", timestamp=None): return ProgressMessage.model_construct( level=level, step=step, text=text, timestamp=timestamp diff --git a/tests/unit/ai_decomp/test_service.py b/tests/unit/ai_decomp/test_service.py index d3ea6ff..63453f6 100644 --- a/tests/unit/ai_decomp/test_service.py +++ b/tests/unit/ai_decomp/test_service.py @@ -1,3 +1,4 @@ +import ast import json import pathlib import threading @@ -750,3 +751,38 @@ def test_view_writing_callbacks_stay_marshalled_onto_the_ui_thread(callback): f"{callback} writes to the Qt view from a stream/poll worker thread, " "so it must keep @execute_ui" ) + + +def _keyword_calls(tree, func_name): + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + target = node.func + name = target.attr if isinstance(target, ast.Attribute) else getattr(target, "id", None) + if name == func_name: + yield node + + +@pytest.mark.parametrize("dialog", ["show_info_dialog", "show_error_dialog"]) +def test_coordinators_name_dialog_arguments_the_way_base_coordinator_declares_them(dialog): + coordinators = pathlib.Path("reai_toolkit/app/coordinators") + base = ast.parse((coordinators / "base_coordinator.py").read_text()) + declared = { + arg.arg + for node in ast.walk(base) + if isinstance(node, ast.FunctionDef) and node.name == dialog + for arg in node.args.args[1:] + } + assert declared + + offenders = [ + f"{path.name}:{call.lineno}" + for path in coordinators.glob("*.py") + for call in _keyword_calls(ast.parse(path.read_text()), dialog) + if any(kw.arg is not None and kw.arg not in declared for kw in call.keywords) + ] + + assert offenders == [], ( + f"BaseCoordinator.{dialog} takes {sorted(declared)}; any other keyword raises " + f"TypeError instead of showing the dialog: {offenders}" + ) From 40d024df6aa0f8f215d6a4b060d5b81d0497b02f Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:57:20 +0100 Subject: [PATCH 07/12] feat: agent chat event changes --- reai_toolkit/app/services/ai_decomp/stream.py | 49 +------ reai_toolkit/app/services/chat/schema.py | 19 ++- reai_toolkit/app/services/chat/sse.py | 57 +++----- reai_toolkit/app/services/sse.py | 50 +++++++ tests/unit/chat/test_sdk_schemas.py | 28 ++++ tests/unit/chat/test_service.py | 6 +- tests/unit/chat/test_sse_parser.py | 133 ++++++++++++++---- 7 files changed, 226 insertions(+), 116 deletions(-) create mode 100644 reai_toolkit/app/services/sse.py diff --git a/reai_toolkit/app/services/ai_decomp/stream.py b/reai_toolkit/app/services/ai_decomp/stream.py index 950aa51..1c650d2 100644 --- a/reai_toolkit/app/services/ai_decomp/stream.py +++ b/reai_toolkit/app/services/ai_decomp/stream.py @@ -4,6 +4,8 @@ from dataclasses import dataclass, field, replace from typing import Callable, Iterable, Iterator, Optional +from reai_toolkit.app.services.sse import iter_sse_frames + EVENT_ATTEMPT_FAILED = "attempt_failed" EVENT_ATTEMPT_STARTED = "attempt_started" EVENT_DECOMP_FAILED = "decomp_failed" @@ -72,53 +74,6 @@ def normalize_event(event_name: Optional[str], payload: dict) -> Optional[Decomp ) -def _iter_lines( - chunks: Iterable[bytes], - stop: Optional[Callable[[], bool]] = None, -) -> Iterator[str]: - buf = b"" - for chunk in chunks: - if stop is not None and stop(): - return - if not chunk: - continue - buf += chunk - parts = buf.split(b"\n") - buf = parts.pop() - for raw in parts: - yield raw.rstrip(b"\r").decode("utf-8", "replace") - if buf: - yield buf.rstrip(b"\r").decode("utf-8", "replace") - - -def iter_sse_frames( - chunks: Iterable[bytes], - stop: Optional[Callable[[], bool]] = None, -) -> Iterator[tuple[Optional[str], str]]: - event_name: Optional[str] = None - data_lines: list[str] = [] - - for line in _iter_lines(chunks, stop=stop): - if not line: - if data_lines: - yield event_name, "\n".join(data_lines) - event_name = None - data_lines = [] - continue - if line.startswith(":"): - continue - name, _, value = line.partition(":") - if value.startswith(" "): - value = value[1:] - if name == "event": - event_name = value - elif name == "data": - data_lines.append(value) - - if data_lines: - yield event_name, "\n".join(data_lines) - - def iter_stream_events( chunks: Iterable[bytes], stop: Optional[Callable[[], bool]] = None, diff --git a/reai_toolkit/app/services/chat/schema.py b/reai_toolkit/app/services/chat/schema.py index 4379f8a..e6a93d9 100644 --- a/reai_toolkit/app/services/chat/schema.py +++ b/reai_toolkit/app/services/chat/schema.py @@ -27,16 +27,21 @@ 14: "RUN_CANCELLED", 15: "CONTEXT_COMPACTED", 16: "TOOL_CONFIRMATION_REQUIRED", - 17: "TOOL_CALL_PROGRESS", + 17: "QUESTION_ASKED", + 18: "USER_ANSWERED", + 19: "TOOL_CALL_PROGRESS", } TERMINAL_EVENTS: frozenset[str] = frozenset( {"RUN_FINISHED", "RUN_ERROR", "RUN_CANCELLED"} ) +ROLE_UNSPECIFIED = 0 +ROLE_SYSTEM = 1 ROLE_USER = 2 -ROLE_SYSTEM = 3 +ROLE_ASSISTANT = 3 ROLE_TOOL = 4 +ROLE_DEVELOPER = 5 @dataclass @@ -133,14 +138,20 @@ def _parse_entity_updates(raw: Any) -> Optional[list[EntityUpdate]]: return out or None -def normalize_event(type_field: Any, leaf: Optional[dict]) -> Optional[ChatEvent]: +def normalize_event( + type_field: Any, + leaf: Optional[dict], + event_name: Optional[str] = None, +) -> Optional[ChatEvent]: """Normalize a raw SSE ``{type, data}`` frame into a :class:`ChatEvent`. ``leaf`` is the nested ``data`` object holding snake_case leaf fields. + ``event_name`` is the SSE ``event:`` field, used when the envelope carries no + resolvable ``type`` — a named event, or a numeric one this table predates. Returns ``None`` for unknown/undecodable event types (mirrors the FE's ``parseApiEvent`` returning null on parse failure). """ - etype = resolve_type(type_field) + etype = resolve_type(type_field) or (event_name or None) if etype is None: return None data: dict = leaf if isinstance(leaf, dict) else {} diff --git a/reai_toolkit/app/services/chat/sse.py b/reai_toolkit/app/services/chat/sse.py index 19644d2..46a3eaa 100644 --- a/reai_toolkit/app/services/chat/sse.py +++ b/reai_toolkit/app/services/chat/sse.py @@ -1,9 +1,9 @@ """Pure Server-Sent-Events frame parsing for the Agent Chat stream. -Ports the SSE read loop of ``Dashboard/utils/v2/agent/agentApi.ts::streamEvents``: -line-based framing, only ``data:`` frames, buffering across chunk boundaries, -``[DONE]`` handling and stopping on terminal events. No Qt / IDA / SDK imports, -so it is unit-testable against a fake byte-chunk iterator. +Reads blank-line delimited frames, so both the envelope's ``type`` and the SSE +``event:`` name are available — the backend transmits both, and omits ``event:`` +only for the default ``message`` name. No Qt / IDA / SDK imports, so it is +unit-testable against a fake byte-chunk iterator. """ from __future__ import annotations @@ -16,30 +16,23 @@ ChatEvent, normalize_event, ) +from reai_toolkit.app.services.sse import iter_sse_frames -def parse_sse_data_line(line: str) -> Optional[dict]: - """Decode a single SSE line into its JSON object, or ``None``. - - Only ``data:`` lines carry payloads; ``event:`` / ``id:`` / comment (``:``) - lines and the ``[DONE]`` sentinel are ignored. - """ - line = line.strip() - if not line.startswith("data:"): - return None - payload = line[len("data:"):].strip() - if not payload or payload == "[DONE]": +def parse_sse_data(data: str) -> Optional[dict]: + data = data.strip() + if not data or data == "[DONE]": return None try: - obj = json.loads(payload) + obj = json.loads(data) except (ValueError, TypeError): return None return obj if isinstance(obj, dict) else None -def event_from_frame(obj: dict) -> Optional[ChatEvent]: +def event_from_frame(obj: dict, event_name: Optional[str] = None) -> Optional[ChatEvent]: """Turn a decoded ``{type, event_id, data}`` envelope into a ChatEvent.""" - ev = normalize_event(obj.get("type"), obj.get("data")) + ev = normalize_event(obj.get("type"), obj.get("data"), event_name) if ev is None: return None eid = obj.get("event_id") @@ -57,23 +50,13 @@ def iter_sse_events( ``stop`` is polled between chunks for cooperative cancellation. Iteration stops after a terminal event (RUN_FINISHED / RUN_ERROR / RUN_CANCELLED). """ - buf = b"" - for chunk in chunks: - if stop is not None and stop(): - return - if not chunk: + for event_name, data in iter_sse_frames(chunks, stop=stop): + obj = parse_sse_data(data) + if obj is None: + continue + ev = event_from_frame(obj, event_name) + if ev is None: continue - buf += chunk - parts = buf.split(b"\n") - buf = parts.pop() - for raw in parts: - line = raw.rstrip(b"\r").decode("utf-8", "replace") - obj = parse_sse_data_line(line) - if obj is None: - continue - ev = event_from_frame(obj) - if ev is None: - continue - yield ev - if ev.type in TERMINAL_EVENTS: - return + yield ev + if ev.type in TERMINAL_EVENTS: + return diff --git a/reai_toolkit/app/services/sse.py b/reai_toolkit/app/services/sse.py new file mode 100644 index 0000000..fd71702 --- /dev/null +++ b/reai_toolkit/app/services/sse.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from typing import Callable, Iterable, Iterator, Optional + + +def iter_lines( + chunks: Iterable[bytes], + stop: Optional[Callable[[], bool]] = None, +) -> Iterator[str]: + buf = b"" + for chunk in chunks: + if stop is not None and stop(): + return + if not chunk: + continue + buf += chunk + parts = buf.split(b"\n") + buf = parts.pop() + for raw in parts: + yield raw.rstrip(b"\r").decode("utf-8", "replace") + if buf: + yield buf.rstrip(b"\r").decode("utf-8", "replace") + + +def iter_sse_frames( + chunks: Iterable[bytes], + stop: Optional[Callable[[], bool]] = None, +) -> Iterator[tuple[Optional[str], str]]: + event_name: Optional[str] = None + data_lines: list[str] = [] + + for line in iter_lines(chunks, stop=stop): + if not line: + if data_lines: + yield event_name, "\n".join(data_lines) + event_name = None + data_lines = [] + continue + if line.startswith(":"): + continue + name, _, value = line.partition(":") + if value.startswith(" "): + value = value[1:] + if name == "event": + event_name = value + elif name == "data": + data_lines.append(value) + + if data_lines: + yield event_name, "\n".join(data_lines) diff --git a/tests/unit/chat/test_sdk_schemas.py b/tests/unit/chat/test_sdk_schemas.py index 3c107a0..e7efd30 100644 --- a/tests/unit/chat/test_sdk_schemas.py +++ b/tests/unit/chat/test_sdk_schemas.py @@ -1,5 +1,6 @@ """Schema-shape assertions: fail loudly if the revengai conversation API drifts.""" +import revengai from revengai import ( ConfirmToolInputBody, Conversation, @@ -11,6 +12,8 @@ ) from revengai.models.event import Event +from reai_toolkit.app.services.chat.schema import EVENT_TYPE_NAMES + def test_conversations_api_has_expected_methods(): for name in ( @@ -51,3 +54,28 @@ def test_conversation_with_events_fields(): def test_event_fields(): assert {"type", "role", "data", "event_id"} <= set(Event.model_fields) + + +def _sdk_agent_event_names() -> set[str]: + names = set() + for attr in dir(revengai): + if not attr.startswith("Event"): + continue + suffix = attr[len("Event"):] + if suffix and suffix.isupper(): + names.add(suffix) + return names + + +def test_every_agent_event_the_sdk_types_can_be_named_by_the_numeric_table(): + known = {name.replace("_", "") for name in EVENT_TYPE_NAMES.values()} + missing = sorted(_sdk_agent_event_names() - known) + + assert missing == [], ( + "the SDK types agent events this plugin cannot name; a numerically typed " + f"replay of one would decode as the wrong event: {missing}" + ) + + +def test_the_replay_type_is_numeric_so_the_table_is_load_bearing(): + assert Event.model_fields["type"].annotation is int diff --git a/tests/unit/chat/test_service.py b/tests/unit/chat/test_service.py index 2988bf2..4d16833 100644 --- a/tests/unit/chat/test_service.py +++ b/tests/unit/chat/test_service.py @@ -108,9 +108,9 @@ def test_get_conversation_replays_events(service, api): SimpleNamespace(type=6, role=2, data={"message_id": "u1"}, event_id=1), SimpleNamespace(type=7, role=2, data={"delta": "why?"}, event_id=2), SimpleNamespace(type=8, role=2, data={}, event_id=3), - SimpleNamespace(type=6, role=1, data={"message_id": "a1"}, event_id=4), - SimpleNamespace(type=7, role=1, data={"delta": "because"}, event_id=5), - SimpleNamespace(type=8, role=1, data={}, event_id=6), + SimpleNamespace(type=6, role=3, data={"message_id": "a1"}, event_id=4), + SimpleNamespace(type=7, role=3, data={"delta": "because"}, event_id=5), + SimpleNamespace(type=8, role=3, data={}, event_id=6), ] api.get_conversation.return_value = SimpleNamespace( conversation_uuid="cid", title="T", events=events diff --git a/tests/unit/chat/test_sse_parser.py b/tests/unit/chat/test_sse_parser.py index 6a4815f..bb4fdc4 100644 --- a/tests/unit/chat/test_sse_parser.py +++ b/tests/unit/chat/test_sse_parser.py @@ -1,31 +1,38 @@ """SSE frame parser tests — feed fake byte chunks, no network.""" -from reai_toolkit.app.services.chat.schema import EntityRef, normalize_event +from reai_toolkit.app.services.chat.schema import ( + EVENT_TYPE_NAMES, + ROLE_ASSISTANT, + ROLE_DEVELOPER, + ROLE_SYSTEM, + ROLE_TOOL, + ROLE_UNSPECIFIED, + ROLE_USER, + EntityRef, + normalize_event, + resolve_type, +) from reai_toolkit.app.services.chat.sse import ( event_from_frame, iter_sse_events, - parse_sse_data_line, + parse_sse_data, ) -def _frame(obj_json: str) -> bytes: - return f"data: {obj_json}\n".encode() - +def _frame(obj_json: str, event: str | None = None) -> bytes: + head = f"event: {event}\n" if event else "" + return f"{head}data: {obj_json}\n\n".encode() -def test_parse_data_line_ignores_non_data(): - assert parse_sse_data_line("event: message") is None - assert parse_sse_data_line(": comment") is None - assert parse_sse_data_line("id: 5") is None - assert parse_sse_data_line("") is None +def test_parse_data_ignores_empty_and_done(): + assert parse_sse_data("") is None + assert parse_sse_data("[DONE]") is None + assert parse_sse_data("{not json}") is None + assert parse_sse_data("[1, 2]") is None -def test_parse_data_line_done_and_empty(): - assert parse_sse_data_line("data: [DONE]") is None - assert parse_sse_data_line("data:") is None - -def test_parse_data_line_decodes_json(): - assert parse_sse_data_line('data: {"type": 1}') == {"type": 1} +def test_parse_data_decodes_json(): + assert parse_sse_data('{"type": 1}') == {"type": 1} def test_event_from_frame_tracks_event_id(): @@ -65,7 +72,7 @@ def test_frames_split_across_chunk_boundaries(): def test_crlf_line_endings(): - frame = b'data: {"type": "TITLE_UPDATED", "data": {"title": "T"}}\r\n' + frame = b'data: {"type": "TITLE_UPDATED", "data": {"title": "T"}}\r\n\r\n' events = list(iter_sse_events([frame])) assert events[0].title == "T" @@ -79,23 +86,55 @@ def test_multiple_frames_in_one_chunk(): assert [e.type for e in events] == ["RUN_STARTED", "TEXT_MESSAGE_CONTENT"] -def test_event_and_id_lines_ignored(): - chunk = b"event: message\nid: 3\n" + _frame('{"type": 1}') +def test_a_final_frame_without_its_blank_line_is_still_read(): + events = list(iter_sse_events([b'data: {"type": 1}\n'])) + assert [e.type for e in events] == ["RUN_STARTED"] + + +def test_comment_and_id_lines_are_skipped(): + chunk = b': keep-alive\nid: 3\ndata: {"type": 1}\n\n' events = list(iter_sse_events([chunk])) assert [e.type for e in events] == ["RUN_STARTED"] +def test_a_named_event_is_read_when_the_envelope_has_no_usable_type(): + chunk = _frame('{"data": {"title": "T"}}', event="TITLE_UPDATED") + events = list(iter_sse_events([chunk])) + assert [e.type for e in events] == ["TITLE_UPDATED"] + assert events[0].title == "T" + + +def test_a_named_event_rescues_a_number_this_table_does_not_know(): + chunk = _frame('{"type": 999, "data": {"title": "T"}}', event="TITLE_UPDATED") + events = list(iter_sse_events([chunk])) + assert [e.type for e in events] == ["TITLE_UPDATED"] + + +def test_the_envelope_type_wins_over_the_event_line(): + chunk = _frame('{"type": 13, "data": {"title": "T"}}', event="RUN_STARTED") + events = list(iter_sse_events([chunk])) + assert [e.type for e in events] == ["TITLE_UPDATED"] + + +def test_a_frame_carrying_only_an_event_line_yields_nothing(): + events = list(iter_sse_events([b"event: RUN_STARTED\n\n"])) + assert events == [] + + +def test_multi_line_data_is_joined_before_decoding(): + chunk = b'data: {"type": 1,\ndata: "event_id": 4}\n\n' + events = list(iter_sse_events([chunk])) + assert [(e.type, e.event_id) for e in events] == [("RUN_STARTED", 4)] + + def test_done_sentinel_and_blank_lines_skipped(): - chunk = b"\n" + _frame('{"type": 1}') + b"data: [DONE]\n" + chunk = b"\n" + _frame('{"type": 1}') + b"data: [DONE]\n\n" events = list(iter_sse_events([chunk])) assert [e.type for e in events] == ["RUN_STARTED"] def test_terminal_event_stops_iteration(): - chunk = ( - _frame('{"type": 2}') - + _frame('{"type": 1}') - ) + chunk = _frame('{"type": 2}') + _frame('{"type": 1}') events = list(iter_sse_events([chunk])) assert [e.type for e in events] == ["RUN_FINISHED"] @@ -112,7 +151,7 @@ def stop(): def test_bad_json_frame_is_skipped(): - chunk = b"data: {not json}\n" + _frame('{"type": 1}') + chunk = _frame("{not json}") + _frame('{"type": 1}') events = list(iter_sse_events([chunk])) assert [e.type for e in events] == ["RUN_STARTED"] @@ -157,3 +196,47 @@ def test_tool_result_entity_refs_parsed(): assert ev.updated[0].refs == [ EntityRef(id=2015699787, name="region_position", vaddr=4198416) ] + + +def test_the_numeric_table_matches_the_persisted_backend_enum(): + assert EVENT_TYPE_NAMES == { + 1: "RUN_STARTED", + 2: "RUN_FINISHED", + 3: "RUN_ERROR", + 4: "STEP_STARTED", + 5: "STEP_FINISHED", + 6: "TEXT_MESSAGE_START", + 7: "TEXT_MESSAGE_CONTENT", + 8: "TEXT_MESSAGE_END", + 9: "TOOL_CALL_START", + 10: "TOOL_CALL_ARGS_DELTA", + 11: "TOOL_CALL_END", + 12: "TOOL_CALL_RESULT", + 13: "TITLE_UPDATED", + 14: "RUN_CANCELLED", + 15: "CONTEXT_COMPACTED", + 16: "TOOL_CONFIRMATION_REQUIRED", + 17: "QUESTION_ASKED", + 18: "USER_ANSWERED", + 19: "TOOL_CALL_PROGRESS", + } + + +def test_replayed_question_asked_is_not_mistaken_for_tool_call_progress(): + assert resolve_type(17) == "QUESTION_ASKED" + assert resolve_type(19) == "TOOL_CALL_PROGRESS" + + +def test_event_type_zero_stays_undecodable(): + assert resolve_type(0) is None + + +def test_the_role_table_matches_the_persisted_backend_enum(): + assert ( + ROLE_UNSPECIFIED, + ROLE_SYSTEM, + ROLE_USER, + ROLE_ASSISTANT, + ROLE_TOOL, + ROLE_DEVELOPER, + ) == (0, 1, 2, 3, 4, 5) From c089e6008be41cc48528fde2c2341594ea2674e7 Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:35:20 +0100 Subject: [PATCH 08/12] feat: base line mapping changes --- .../app/components/tabs/ai_decomp_tab.py | 36 +++++ .../app/coordinators/ai_decomp_coordinator.py | 70 +++++++++- .../app/coordinators/ai_decomp_render.py | 19 +++ .../services/ai_decomp/ai_decomp_service.py | 80 +++++++++++ .../app/services/ai_decomp/attribution.py | 121 ++++++++++++++++ reai_toolkit/hooks/reactive.py | 57 ++++++++ tests/idalib/ai_decomp/test_smoke.py | 68 +++++++++ tests/ui/ai_decomp_ui_runner.py | 43 ++++++ tests/ui/test_gui_ai_decomp.py | 4 + tests/unit/ai_decomp/test_attribution.py | 132 ++++++++++++++++++ tests/unit/ai_decomp/test_render.py | 47 +++++++ tests/unit/ai_decomp/test_sdk_schemas.py | 30 +++- tests/unit/ai_decomp/test_service.py | 114 ++++++++++++++- 13 files changed, 818 insertions(+), 3 deletions(-) create mode 100644 reai_toolkit/app/services/ai_decomp/attribution.py create mode 100644 tests/unit/ai_decomp/test_attribution.py diff --git a/reai_toolkit/app/components/tabs/ai_decomp_tab.py b/reai_toolkit/app/components/tabs/ai_decomp_tab.py index 80bd836..7938678 100644 --- a/reai_toolkit/app/components/tabs/ai_decomp_tab.py +++ b/reai_toolkit/app/components/tabs/ai_decomp_tab.py @@ -15,6 +15,12 @@ _IDENTIFIER_RE = re.compile(r"[A-Za-z_]\w*\Z") +_FULL_WIDTH_SELECTION = getattr( + getattr(QtGui.QTextFormat, "Property", QtGui.QTextFormat), "FullWidthSelection" +) + +ATTRIBUTION_HIGHLIGHT_RGB = "#3d4a63" + def _is_identifier(word: str) -> bool: return bool(word) and _IDENTIFIER_RE.match(word) is not None @@ -93,6 +99,7 @@ def __init__(self, on_closed: Optional[Callable[[], None]] = None) -> None: self.on_rate_up: Callable[[], None] | None = None self.on_rate_down: Callable[[], None] | None = None self.on_use_predicted_name: Callable[[str], None] | None = None + self.on_line_focus: Callable[[int], None] | None = None self._parent_window: QtWidgets.QWidget | None = None self._editor: _DecompEditor | None = None self._refresh_btn: QtWidgets.QPushButton | None = None @@ -184,6 +191,7 @@ def OnCreate(self, form) -> None: self._editor.renameRequested.connect(self._on_rename_requested) self._editor.commentEditRequested.connect(self._on_edit_comment_requested) self._editor.commentRemoveRequested.connect(self._on_remove_comment_requested) + self._editor.cursorPositionChanged.connect(self._on_cursor_moved) # Monospace font tuned for IDA font = QtGui.QFont( @@ -268,6 +276,34 @@ def _on_remove_comment_requested(self, line: int) -> None: if self.on_remove_comment: self.on_remove_comment(line) + def _on_cursor_moved(self) -> None: + if self._editor and self.on_line_focus: + self.on_line_focus(self._editor.textCursor().blockNumber()) + + @execute_ui + def set_highlighted_lines(self, lines) -> None: + if not self._editor: + return + + fmt = QtGui.QTextCharFormat() + fmt.setBackground(QtGui.QColor(ATTRIBUTION_HIGHLIGHT_RGB)) + fmt.setProperty(_FULL_WIDTH_SELECTION, True) + + document = self._editor.document() + selections = [] + for line in sorted(lines): + block = document.findBlockByNumber(line) + if not block.isValid(): + continue + selection = QtWidgets.QTextEdit.ExtraSelection() + selection.format = fmt + cursor = QtGui.QTextCursor(block) + cursor.clearSelection() + selection.cursor = cursor + selections.append(selection) + + self._editor.setExtraSelections(selections) + # --- public API ------------------------------------------------ @execute_ui def update_view_content(self, code: str, follow_tail: bool = False) -> None: diff --git a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py index 45dfbd4..f1d9573 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py +++ b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py @@ -15,17 +15,20 @@ from reai_toolkit.app.coordinators.ai_decomp_render import ( RENAME_NOT_READY, RenderModel, + display_rows_for_source_lines, render_progress, render_stream, render_view_with_map, resolve_rename_target, + source_line_at, ) from reai_toolkit.app.coordinators.base_coordinator import BaseCoordinator from reai_toolkit.app.core.shared_schema import GenericApiReturn from reai_toolkit.app.factory import DialogFactory from reai_toolkit.app.services.ai_decomp.ai_decomp_service import AiDecompService +from reai_toolkit.app.services.ai_decomp.attribution import AttributionMap from reai_toolkit.app.services.ai_decomp.stream import StreamState -from reai_toolkit.hooks.reactive import AiDecompFunctionViewHooks +from reai_toolkit.hooks.reactive import AiDecompFunctionViewHooks, LineAttributionHooks class AiDecompCoordinator(BaseCoordinator): @@ -46,6 +49,10 @@ def __init__( self._current_summary: SummaryData | None = None self._current_comments: CommentsData | None = None self._current_tokenised: GetTokensResponse | None = None + self._current_attributions: AttributionMap | None = None + self._attribution_hooks: LineAttributionHooks | None = None + self._lit_addresses: frozenset[int] = frozenset() + self._lit_rows: tuple[int, ...] = () self._baseline: RenderModel | None = None def enable_function_tracking(self) -> None: @@ -76,15 +83,21 @@ def run_dialog(self) -> None: self._decomp_view.on_rate_up = self.rate_up self._decomp_view.on_rate_down = self.rate_down self._decomp_view.on_use_predicted_name = self.apply_predicted_name + self._decomp_view.on_line_focus = self.focus_decomp_line self._decomp_view.Create(self._decomp_view.TITLE) + if self._attribution_hooks is None: + self._attribution_hooks = LineAttributionHooks(self) + self._attribution_hooks.hook() def start_decompilation(self, ea: int) -> None: self._current_decomp = None self._current_summary = None self._current_comments = None self._current_tokenised = None + self._current_attributions = None self._baseline = None self._current_func_vaddr = ea + self._clear_attribution_highlights() self.run_dialog() @@ -128,6 +141,7 @@ def _dispatch_task(self, ea: int) -> None: on_tokenised=lambda response: self._on_tokenised_complete(ea, response), on_progress=lambda progress: self._on_progress(ea, progress), on_stream=lambda state: self._on_stream(ea, state), + on_attributions=lambda mapping: self._on_attributions(ea, mapping), ) @execute_ui @@ -218,6 +232,52 @@ def _on_tokenised_complete( return self._current_tokenised = response.data + @execute_ui + def _on_attributions(self, ea: int, mapping: AttributionMap) -> None: + if ea != self._current_func_vaddr: + return + self._current_attributions = None if mapping.is_empty() else mapping + + def focus_decomp_line(self, display_line: int) -> None: + if self._current_attributions is None or self._baseline is None: + return + source_line = source_line_at(self._baseline, display_line) + addresses = ( + () + if source_line is None + else self._current_attributions.addresses_for_decomp_line(source_line - 1) + ) + self._light_disassembly(addresses) + + def on_disassembly_ea(self, ea: int) -> None: + if self._current_attributions is None or self._baseline is None: + return + decomp_lines = self._current_attributions.decomp_lines_for_address(ea) + self._light_decomp( + display_rows_for_source_lines( + self._baseline, [line + 1 for line in decomp_lines] + ) + ) + + def _light_disassembly(self, addresses) -> None: + wanted = frozenset(addresses) + if self._attribution_hooks is None or wanted == self._lit_addresses: + return + self._lit_addresses = wanted + self._attribution_hooks.set_addresses(wanted) + self.refresh_disassembly_view() + + def _light_decomp(self, rows) -> None: + wanted = tuple(rows) + if self._decomp_view is None or wanted == self._lit_rows: + return + self._lit_rows = wanted + self._decomp_view.set_highlighted_lines(wanted) + + def _clear_attribution_highlights(self) -> None: + self._light_disassembly(()) + self._light_decomp(()) + def refresh_current(self) -> None: ea = self._current_func_vaddr if ea is None: @@ -375,8 +435,16 @@ def _rerender(self) -> None: comments=self._current_comments, ) self._decomp_view.update_view_content(rendered) + self._lit_rows = () def _on_pane_closed(self) -> None: self._decomp_view = None + if self._attribution_hooks is not None: + self._attribution_hooks.unhook() + self._attribution_hooks = None + self.refresh_disassembly_view() + self._current_attributions = None + self._lit_addresses = frozenset() + self._lit_rows = () self.disable_function_tracking() self.log.debug("AI Decomp view closed, reference cleared.") diff --git a/reai_toolkit/app/coordinators/ai_decomp_render.py b/reai_toolkit/app/coordinators/ai_decomp_render.py index 5d4105a..d5fd48f 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_render.py +++ b/reai_toolkit/app/coordinators/ai_decomp_render.py @@ -222,6 +222,25 @@ def index_of_identifier(line: str, word: str) -> int: return idents.index(word) if word in idents else -1 +def source_line_at(model: RenderModel, display_line: int) -> Optional[int]: + if not (0 <= display_line < len(model.display_source)): + return None + if not model.display_is_code[display_line]: + return None + return model.display_source[display_line] + + +def display_rows_for_source_lines(model: RenderModel, source_lines) -> list[int]: + wanted = set(source_lines) + if not wanted: + return [] + return [ + row + for row, source in enumerate(model.display_source) + if source in wanted and model.display_is_code[row] + ] + + def effective_values(tokens: "GetTokensResponse") -> dict[str, str]: rendered = tokens.placeholder_to_rendered_token or {} overrides = tokens.placeholder_to_user_override or {} diff --git a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py index 0929239..1560452 100644 --- a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py +++ b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py @@ -9,6 +9,7 @@ Configuration, FunctionMapping, FunctionsAIDecompilationApi, + FunctionsCoreApi, ) from revengai.models.ai_decompilation_rating import AiDecompilationRating from revengai.models.comments_data import CommentsData @@ -28,6 +29,11 @@ from reai_toolkit.app.core.shared_schema import GenericApiReturn from reai_toolkit.app.core.utils import parse_exception from reai_toolkit.app.interfaces.thread_service import IThreadService +from reai_toolkit.app.services.ai_decomp.attribution import ( + AttributionMap, + asm_row_addresses, + invert_attributions, +) from reai_toolkit.app.services.ai_decomp.stream import StreamState, iter_stream_events @@ -47,6 +53,7 @@ def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration) self._summary_cache: dict[int, SummaryData] = {} self._comments_cache: dict[int, CommentsData] = {} self._tokenised_cache: dict[int, GetTokensResponse] = {} + self._attribution_cache: dict[int, AttributionMap] = {} self._inflight: dict[int, threading.Event] = {} self._inflight_lock = threading.Lock() self._active_streams: dict[int, urllib3.HTTPResponse] = {} @@ -70,6 +77,7 @@ def stop_worker(self) -> None: self._summary_cache.clear() self._comments_cache.clear() self._tokenised_cache.clear() + self._attribution_cache.clear() def peek_decomp(self, ea: int) -> DecompilationData | None: function_id = self._get_function_id(start_ea=ea) @@ -85,6 +93,7 @@ def invalidate(self, function_id: int) -> None: self._summary_cache.pop(function_id, None) self._comments_cache.pop(function_id, None) self._tokenised_cache.pop(function_id, None) + self._attribution_cache.pop(function_id, None) with self._inflight_lock: evt = self._inflight.pop(function_id, None) if evt is not None: @@ -105,6 +114,7 @@ def start_ai_decomp_task( on_tokenised: Callable[[GenericApiReturn[GetTokensResponse]], None] | None = None, on_progress: Callable[[WorkflowProgress], None] | None = None, on_stream: Callable[[StreamState], None] | None = None, + on_attributions: Callable[[AttributionMap], None] | None = None, ) -> None: function_id = self._get_function_id(start_ea=ea) if function_id is None: @@ -128,6 +138,7 @@ def start_ai_decomp_task( on_tokenised, on_progress, on_stream, + on_attributions, ), name=f"reai-aidecomp-{function_id}", daemon=True, @@ -375,6 +386,7 @@ def _run_task( on_tokenised: Callable[[GenericApiReturn[GetTokensResponse]], None] | None = None, on_progress: Callable[[WorkflowProgress], None] | None = None, on_stream: Callable[[StreamState], None] | None = None, + on_attributions: Callable[[AttributionMap], None] | None = None, ) -> None: try: if stop_event.is_set(): @@ -388,6 +400,8 @@ def _run_task( self._run_comments_phase(function_id, stop_event, on_comments) if on_tokenised is not None: self._run_tokenised_phase(function_id, stop_event, on_tokenised) + if on_attributions is not None: + self._run_attribution_phase(function_id, stop_event, on_attributions) except Exception as e: logger.error(f"RevEng.AI: AI decompilation task crashed for {function_id}: {e}") finally: @@ -927,6 +941,72 @@ def _run_tokenised_phase( ), ) + def _fetch_line_attributions(self, function_id: int) -> tuple[Any, str | None]: + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + data = FunctionsAIDecompilationApi( + api_client + ).v3_get_ai_decompilation_line_attributions(function_id=function_id) + return data, None + except ApiException as e: + return None, _format_api_error(e) + except Exception as e: + return None, f"Unexpected error fetching line attributions: {e}" + + def _fetch_blocks(self, function_id: int) -> tuple[Any, str | None]: + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + data = FunctionsCoreApi(api_client).get_function_blocks_0( + function_id=function_id + ) + return data, None + except ApiException as e: + return None, _format_api_error(e) + except Exception as e: + return None, f"Unexpected error fetching function blocks: {e}" + + def _run_attribution_phase( + self, + function_id: int, + stop_event: threading.Event, + on_attributions: Callable[[AttributionMap], None], + ) -> None: + cached = self._attribution_cache.get(function_id) + if cached is not None: + self._safe_dispatch(stop_event, on_attributions, cached) + return + + data, err = self._fetch_line_attributions(function_id) + if data is None: + logger.debug( + f"RevEng.AI: no line attributions for function {function_id}: {err}" + ) + return + + attributions = invert_attributions( + getattr( + data, + "disassembly_line_number_to_ai_decompilation_line_numbers", + None, + ) + ) + if attributions.is_empty() or stop_event.is_set(): + self._attribution_cache[function_id] = AttributionMap() + return + + blocks, blocks_err = self._fetch_blocks(function_id) + if blocks is None: + logger.debug( + f"RevEng.AI: no disassembly blocks for function {function_id}: {blocks_err}" + ) + return + + mapping = AttributionMap( + attributions, asm_row_addresses(getattr(blocks, "basic_blocks", None)) + ) + self._attribution_cache[function_id] = mapping + self._safe_dispatch(stop_event, on_attributions, mapping) + def _poll_workflow( self, function_id: int, diff --git a/reai_toolkit/app/services/ai_decomp/attribution.py b/reai_toolkit/app/services/ai_decomp/attribution.py new file mode 100644 index 0000000..2256760 --- /dev/null +++ b/reai_toolkit/app/services/ai_decomp/attribution.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from typing import Any, Optional + +DISASSEMBLY_ROW_OFFSET = 1 + +_LINE_ADDRESS_RE = re.compile(r"^\s*(0[xX][0-9a-fA-F]+)") + + +def _as_index(value: Any) -> Optional[int]: + if isinstance(value, bool): + return None + if isinstance(value, int): + return value if value >= 0 else None + if isinstance(value, str): + try: + parsed = int(value, 10) + except ValueError: + return None + return parsed if parsed >= 0 else None + return None + + +@dataclass +class LineAttributions: + disasm_to_decomp: dict[int, list[int]] = field(default_factory=dict) + decomp_to_disasm: dict[int, list[int]] = field(default_factory=dict) + + def is_empty(self) -> bool: + return not self.disasm_to_decomp + + +def invert_attributions(forward: Any) -> LineAttributions: + if not isinstance(forward, dict): + return LineAttributions() + + disasm_to_decomp: dict[int, list[int]] = {} + decomp_to_disasm: dict[int, list[int]] = {} + + for raw_line, raw_decomp_lines in forward.items(): + disasm_line = _as_index(raw_line) + if disasm_line is None or not raw_decomp_lines: + continue + decomp_lines = [ + line + for line in (_as_index(value) for value in raw_decomp_lines) + if line is not None + ] + if not decomp_lines: + continue + disasm_to_decomp[disasm_line] = decomp_lines + for decomp_line in decomp_lines: + decomp_to_disasm.setdefault(decomp_line, []).append(disasm_line) + + for disasm_lines in decomp_to_disasm.values(): + disasm_lines.sort() + + return LineAttributions(disasm_to_decomp, decomp_to_disasm) + + +def _block_min_addr(block: Any) -> int: + value = block.get("min_addr") if isinstance(block, dict) else None + return value if isinstance(value, int) and not isinstance(value, bool) else 0 + + +def _block_asm(block: Any) -> list[str]: + value = block.get("asm") if isinstance(block, dict) else None + if not isinstance(value, list): + return [] + return [line for line in value if isinstance(line, str)] + + +def _line_address(line: str) -> Optional[int]: + match = _LINE_ADDRESS_RE.match(line) + return int(match.group(1), 16) if match is not None else None + + +def asm_row_addresses(basic_blocks: Any) -> list[Optional[int]]: + if not isinstance(basic_blocks, list): + return [] + rows: list[Optional[int]] = [] + for block in sorted(basic_blocks, key=_block_min_addr): + for line in _block_asm(block): + rows.append(_line_address(line)) + return rows + + +class AttributionMap: + def __init__( + self, + attributions: Optional[LineAttributions] = None, + row_addresses: Optional[list[Optional[int]]] = None, + ) -> None: + self.attributions = attributions or LineAttributions() + self.row_addresses = row_addresses or [] + self._rows_by_address: dict[int, list[int]] = {} + for row, address in enumerate(self.row_addresses): + if address is not None: + self._rows_by_address.setdefault(address, []).append(row) + + def is_empty(self) -> bool: + return self.attributions.is_empty() or not self._rows_by_address + + def addresses_for_decomp_line(self, decomp_line: int) -> list[int]: + addresses: set[int] = set() + for disasm_line in self.attributions.decomp_to_disasm.get(decomp_line, ()): + row = disasm_line + DISASSEMBLY_ROW_OFFSET + if 0 <= row < len(self.row_addresses): + address = self.row_addresses[row] + if address is not None: + addresses.add(address) + return sorted(addresses) + + def decomp_lines_for_address(self, address: int) -> list[int]: + decomp_lines: set[int] = set() + for row in self._rows_by_address.get(address, ()): + disasm_line = row - DISASSEMBLY_ROW_OFFSET + decomp_lines.update(self.attributions.disasm_to_decomp.get(disasm_line, ())) + return sorted(decomp_lines) diff --git a/reai_toolkit/hooks/reactive.py b/reai_toolkit/hooks/reactive.py index 860a02c..7de505b 100644 --- a/reai_toolkit/hooks/reactive.py +++ b/reai_toolkit/hooks/reactive.py @@ -66,6 +66,63 @@ def screen_ea_changed(self, ea: int, prev_ea: int) -> None: logger.error(f"[ChatContextHooks] context update failed: {e}") +ATTRIBUTION_BG_COLOR = 0xFF000000 | 0x63_4A_3D + + +class LineAttributionHooks(kw.UI_Hooks): + """ + Paints the disassembly lines a decompiled line was attributed to, and + reports cursor movement in the disassembly so the panel can light the + matching decompilation lines. + + Rendering is transient: nothing is written to the IDB. + """ + + def __init__(self, coordinator): + super().__init__() + self.coordinator = coordinator + self._addresses: frozenset[int] = frozenset() + self._is_hooked = False + + def hook(self) -> bool: + if self._is_hooked: + return False + self._is_hooked = super().hook() + return self._is_hooked + + def unhook(self) -> None: + if self._is_hooked: + super().unhook() + self._is_hooked = False + self._addresses = frozenset() + + def set_addresses(self, addresses) -> None: + self._addresses = frozenset(addresses) + + def has_addresses(self) -> bool: + return bool(self._addresses) + + def get_lines_rendering_info(self, out, widget, rin) -> None: + if not self._addresses: + return + if kw.get_widget_type(widget) != idaapi.BWN_DISASM: + return + for section in rin.sections_lines: + for line in section: + if line.at.toea() in self._addresses: + out.entries.push_back( + kw.line_rendering_output_entry_t( + line, kw.LROEF_FULL_LINE, ATTRIBUTION_BG_COLOR + ) + ) + + def screen_ea_changed(self, ea: int, prev_ea: int) -> None: + try: + self.coordinator.on_disassembly_ea(ea) + except Exception as e: + logger.debug(f"[LineAttributionHooks] reverse highlight failed: {e}") + + class AiDecompFunctionViewHooks(kw.UI_Hooks): """ Hook that tracks when the screen EA changes (user moves between functions), diff --git a/tests/idalib/ai_decomp/test_smoke.py b/tests/idalib/ai_decomp/test_smoke.py index cb860d9..6455bab 100644 --- a/tests/idalib/ai_decomp/test_smoke.py +++ b/tests/idalib/ai_decomp/test_smoke.py @@ -132,3 +132,71 @@ def test_mutations_and_refresh_under_idalib(loaded_binary, mocker): assert 7 not in service._decomp_cache assert 7 not in service._tokenised_cache assert 7 not in service._comments_cache + + +def test_attribution_highlighting_primitives_under_idalib(loaded_binary): + import ida_kernwin + + from reai_toolkit.hooks.reactive import ATTRIBUTION_BG_COLOR, LineAttributionHooks + + entry = ida_kernwin.line_rendering_output_entry_t( + ida_kernwin.twinline_t(), ida_kernwin.LROEF_FULL_LINE, ATTRIBUTION_BG_COLOR + ) + assert entry.is_bg_color_direct(), "a colour key would follow the theme, not our tint" + assert entry.flags == ida_kernwin.LROEF_FULL_LINE + + hooks = LineAttributionHooks(MagicMock()) + assert not hooks.has_addresses() + hooks.set_addresses([0x1000, 0x1000, 0x1004]) + assert hooks.has_addresses() + hooks.unhook() + assert not hooks.has_addresses() + + +def test_attribution_phase_resolves_real_addresses_under_idalib(loaded_binary, mocker): + ea = next(iter(idautils.Functions()), None) + assert ea is not None + + mocker.patch.object(svc_mod.AiDecompService, "yield_api_client") + api_inst = MagicMock() + mocker.patch.object(svc_mod, "FunctionsAIDecompilationApi", return_value=api_inst) + core_inst = MagicMock() + mocker.patch.object(svc_mod, "FunctionsCoreApi", return_value=core_inst) + + api_inst.get_ai_decompilation.return_value = DecompilationData.model_construct( + status=TaskStatus.COMPLETED.value, decompilation="int sub() { return 0; }" + ) + api_inst.get_ai_decompilation_summary.return_value = SummaryData.model_construct( + ai_summary="", summary="", task_status=TaskStatus.COMPLETED.value + ) + api_inst.get_ai_decompilation_inline_comments.return_value = ( + CommentsData.model_construct( + inline_comments=[], task_status=TaskStatus.COMPLETED.value + ) + ) + api_inst.v3_get_ai_decompilation_line_attributions.return_value = MagicMock( + disassembly_line_number_to_ai_decompilation_line_numbers={"0": [1]} + ) + core_inst.get_function_blocks_0.return_value = MagicMock( + basic_blocks=[ + {"min_addr": ea, "asm": [f"{ea:#x} first", f"{ea + 1:#x} second"]} + ] + ) + + netstore = MagicMock() + netstore.get_function_mapping.return_value.inverse_function_map = {str(ea): 7} + service = AiDecompService(netstore_service=netstore, sdk_config=MagicMock()) + + on_attributions = MagicMock() + service.start_ai_decomp_task( + ea=ea, + on_decomp=MagicMock(), + on_summary=MagicMock(), + on_comments=MagicMock(), + on_attributions=on_attributions, + ) + _wait_mock(on_attributions) + + mapping = on_attributions.call_args[0][0] + assert mapping.addresses_for_decomp_line(1) == [ea + 1] + assert mapping.decomp_lines_for_address(ea + 1) == [1] diff --git a/tests/ui/ai_decomp_ui_runner.py b/tests/ui/ai_decomp_ui_runner.py index 0116698..6810eca 100644 --- a/tests/ui/ai_decomp_ui_runner.py +++ b/tests/ui/ai_decomp_ui_runner.py @@ -51,6 +51,10 @@ def _run(report: dict) -> None: ) from reai_toolkit.app.core.qt_compat import QtWidgets from reai_toolkit.app.core.shared_schema import GenericApiReturn + from reai_toolkit.app.services.ai_decomp.attribution import ( + AttributionMap, + invert_attributions, + ) answers = {"str": "count", "text": "hello"} ida_kernwin.ask_str = lambda default, hist, prompt: answers["str"] @@ -255,6 +259,41 @@ def code_line_row(needle: str) -> int: service.remove_comment.call_args.kwargs.get("line") == 2 ) + seed_plain() + coord._current_attributions = AttributionMap( + invert_attributions({"0": [1]}), [ea + 1, ea, ea + 2] + ) + highlighted: list = [] + coord._attribution_hooks = SimpleNamespace( + set_addresses=lambda addrs: highlighted.append(sorted(addrs)), + has_addresses=lambda: bool(highlighted and highlighted[-1]), + ) + + coord.focus_decomp_line(code_line_row("int v5")) + pump() + report["decomp_line_lights_its_disassembly"] = highlighted[-1:] == [[ea]] + + coord.focus_decomp_line(code_line_row("return v5")) + pump() + report["a_line_with_no_counterpart_lights_nothing"] = highlighted[-1:] == [[]] + + coord.on_disassembly_ea(ea) + pump() + report["disassembly_ea_lights_its_decomp_line"] = ( + len(view._editor.extraSelections()) == 1 + and view._editor.extraSelections()[0].cursor.blockNumber() + == code_line_row("int v5") + ) + + coord.on_disassembly_ea(0xDEAD) + pump() + report["an_unattributed_address_lights_nothing"] = ( + view._editor.extraSelections() == [] + ) + + coord._attribution_hooks = None + coord._current_attributions = None + service.reset_mock() service.peek_decomp.return_value = None view._refresh_btn.click() @@ -290,6 +329,10 @@ def main() -> None: "comment_edit_empty_removes": False, "comment_remove_deletes": False, "comment_remove_args_correct": False, + "decomp_line_lights_its_disassembly": False, + "a_line_with_no_counterpart_lights_nothing": False, + "disassembly_ea_lights_its_decomp_line": False, + "an_unattributed_address_lights_nothing": False, "refresh_button_invalidates": False, } ida_auto.auto_wait() diff --git a/tests/ui/test_gui_ai_decomp.py b/tests/ui/test_gui_ai_decomp.py index 8f2d592..6a22505 100644 --- a/tests/ui/test_gui_ai_decomp.py +++ b/tests/ui/test_gui_ai_decomp.py @@ -32,6 +32,10 @@ "comment_edit_empty_removes", "comment_remove_deletes", "comment_remove_args_correct", + "decomp_line_lights_its_disassembly", + "a_line_with_no_counterpart_lights_nothing", + "disassembly_ea_lights_its_decomp_line", + "an_unattributed_address_lights_nothing", "refresh_button_invalidates", ] diff --git a/tests/unit/ai_decomp/test_attribution.py b/tests/unit/ai_decomp/test_attribution.py new file mode 100644 index 0000000..58b3424 --- /dev/null +++ b/tests/unit/ai_decomp/test_attribution.py @@ -0,0 +1,132 @@ +import pytest + +from reai_toolkit.app.services.ai_decomp.attribution import ( + DISASSEMBLY_ROW_OFFSET, + AttributionMap, + asm_row_addresses, + invert_attributions, +) + + +FORWARD = {"12": [3, 4, 6], "17": [4]} + + +def _blocks(*specs): + return [{"min_addr": addr, "asm": list(asm)} for addr, asm in specs] + + +def test_invert_keeps_the_forward_map_and_builds_the_reverse(): + attributions = invert_attributions(FORWARD) + + assert attributions.disasm_to_decomp == {12: [3, 4, 6], 17: [4]} + assert attributions.decomp_to_disasm == {3: [12], 4: [12, 17], 6: [12]} + + +def test_invert_sorts_the_reverse_lists(): + attributions = invert_attributions({"9": [1], "2": [1], "5": [1]}) + + assert attributions.decomp_to_disasm[1] == [2, 5, 9] + + +def test_invert_drops_lines_with_no_counterpart(): + attributions = invert_attributions({"1": None, "2": [], "3": [7]}) + + assert attributions.disasm_to_decomp == {3: [7]} + + +def test_invert_drops_unusable_keys_and_values(): + attributions = invert_attributions( + {"nope": [1], "-2": [1], "4": ["x", None, 2], "5": [True]} + ) + + assert attributions.disasm_to_decomp == {4: [2]} + + +@pytest.mark.parametrize("forward", [None, [], "", 0]) +def test_invert_of_a_non_mapping_is_empty(forward): + assert invert_attributions(forward).is_empty() + + +def test_an_empty_correspondence_is_ordinary_not_an_error(): + assert invert_attributions({}).is_empty() + + +def test_rows_flatten_blocks_in_address_order(): + blocks = _blocks( + (0x2000, ["0x2000 ret"]), + (0x1000, ["0x1000 push rbp", "0x1004 mov rbp, rsp"]), + ) + + assert asm_row_addresses(blocks) == [0x1000, 0x1004, 0x2000] + + +def test_rows_keep_their_index_when_a_line_has_no_address(): + blocks = _blocks((0x1000, ["; a comment", "0x1000 ret"])) + + assert asm_row_addresses(blocks) == [None, 0x1000] + + +def test_rows_read_addresses_beyond_32_bits(): + blocks = _blocks((0x7FFF00001000, ["0x7fff00001000 ret"])) + + assert asm_row_addresses(blocks) == [0x7FFF00001000] + + +@pytest.mark.parametrize("blocks", [None, {}, "", 0]) +def test_rows_of_an_unusable_blocks_payload_are_empty(blocks): + assert asm_row_addresses(blocks) == [] + + +def test_rows_tolerate_a_block_missing_its_fields(): + blocks = [{}, {"asm": None}, {"min_addr": 1, "asm": ["0x1 nop", 5]}] + + assert asm_row_addresses(blocks) == [0x1] + + +def _map(): + rows = [0x1000 + 4 * i for i in range(20)] + return AttributionMap(invert_attributions(FORWARD), rows) + + +def test_a_decomp_line_resolves_to_its_disassembly_addresses(): + assert _map().addresses_for_decomp_line(4) == [0x1000 + 4 * 13, 0x1000 + 4 * 18] + + +def test_an_address_resolves_back_to_its_decomp_lines(): + assert _map().decomp_lines_for_address(0x1000 + 4 * 13) == [3, 4, 6] + assert _map().decomp_lines_for_address(0x1000 + 4 * 18) == [4] + + +def test_the_disassembly_row_offset_is_applied_in_both_directions(): + rows = [0x1000 + 4 * i for i in range(20)] + amap = AttributionMap(invert_attributions({"12": [3]}), rows) + expected = rows[12 + DISASSEMBLY_ROW_OFFSET] + + assert amap.addresses_for_decomp_line(3) == [expected] + assert amap.decomp_lines_for_address(expected) == [3] + assert amap.decomp_lines_for_address(rows[12]) == [] + + +def test_a_decomp_line_with_no_counterpart_resolves_to_nothing(): + assert _map().addresses_for_decomp_line(99) == [] + assert _map().decomp_lines_for_address(0xDEAD) == [] + + +def test_rows_beyond_the_disassembly_are_skipped_rather_than_wrapping(): + amap = AttributionMap(invert_attributions({"12": [3]}), [0x1000, 0x1004]) + + assert amap.addresses_for_decomp_line(3) == [] + + +def test_one_address_on_several_rows_unions_their_decomp_lines(): + rows = [0x1000] * 20 + amap = AttributionMap(invert_attributions(FORWARD), rows) + + assert amap.decomp_lines_for_address(0x1000) == [3, 4, 6] + + +def test_a_map_without_attributions_or_rows_is_empty(): + assert AttributionMap().is_empty() + assert AttributionMap(invert_attributions(FORWARD), []).is_empty() + assert AttributionMap(invert_attributions({}), [0x1000]).is_empty() + assert not _map().is_empty() diff --git a/tests/unit/ai_decomp/test_render.py b/tests/unit/ai_decomp/test_render.py index 4fae7af..63cdede 100644 --- a/tests/unit/ai_decomp/test_render.py +++ b/tests/unit/ai_decomp/test_render.py @@ -17,6 +17,7 @@ RENAME_NOT_DECOMP_LINE, RENAME_NOT_ON_LINE, RENAME_UNRESOLVED, + display_rows_for_source_lines, effective_values, find_token, index_of_identifier, @@ -27,6 +28,7 @@ render_view_with_map, resolve_rename_target, resolve_token, + source_line_at, ) from reai_toolkit.app.services.ai_decomp.stream import PROSE_TAIL, StreamState @@ -289,6 +291,51 @@ def test_every_declined_rename_carries_a_reason(): assert all(t.placeholder is None and t.reason for t in declines) +def test_source_line_at_maps_a_code_row_back_to_its_decompilation_line(): + model = _model(summary=_summary("S."), comments=_comments([(2, "note")])) + + assert source_line_at(model, 3) == 1 + assert source_line_at(model, 5) == 2 + + +@pytest.mark.parametrize("row", [-1, 0, 4, 99]) +def test_source_line_at_declines_rows_that_are_not_code(row): + model = _model(summary=_summary("S."), comments=_comments([(2, "note")])) + + assert source_line_at(model, row) is None + + +def test_display_rows_finds_the_code_rows_for_source_lines(): + model = _model(summary=_summary("S."), comments=_comments([(2, "note")])) + + assert display_rows_for_source_lines(model, [1, 2]) == [3, 5] + + +def test_display_rows_skips_the_comment_row_that_shares_a_source_line(): + model = _model(comments=_comments([(2, "note")])) + rows = display_rows_for_source_lines(model, [2]) + + assert rows == [2] + assert model.display_is_code[1] is False + assert model.display_source[1] == 2 + + +def test_display_rows_of_nothing_is_nothing(): + model = _model() + + assert display_rows_for_source_lines(model, []) == [] + assert display_rows_for_source_lines(model, [99]) == [] + + +def test_source_line_and_display_rows_round_trip(): + model = _model(summary=_summary("S."), comments=_comments([(2, "note")])) + rows = [row for row, code in enumerate(model.display_is_code) if code] + + for row in rows: + source = source_line_at(model, row) + assert display_rows_for_source_lines(model, [source]) == [row] + + def _pm(text, level="INFO", step="DECOMPILING", timestamp=None): return ProgressMessage.model_construct( level=level, step=step, text=text, timestamp=timestamp diff --git a/tests/unit/ai_decomp/test_sdk_schemas.py b/tests/unit/ai_decomp/test_sdk_schemas.py index d70cb6f..2fd51a5 100644 --- a/tests/unit/ai_decomp/test_sdk_schemas.py +++ b/tests/unit/ai_decomp/test_sdk_schemas.py @@ -4,9 +4,11 @@ import inspect -from revengai import FunctionsAIDecompilationApi +from revengai import FunctionsAIDecompilationApi, FunctionsCoreApi from revengai.models.comments_data import CommentsData from revengai.models.create_ai_decomp_output_body import CreateAIDecompOutputBody +from revengai.models.disassembly_output_body import DisassemblyOutputBody +from revengai.models.line_attributions_data import LineAttributionsData from revengai.models.get_tokens_response import GetTokensResponse from revengai.models.inline_comment import InlineComment from revengai.models.patch_comment_body import PatchCommentBody @@ -35,10 +37,36 @@ def test_ai_decomp_api_exposes_the_methods_the_plugin_calls(): "upsert_ai_decompilation_rating", "v3_get_ai_decompilation_tokens", "v3_upsert_ai_decompilation_overrides", + "v3_get_ai_decompilation_line_attributions", ): assert callable(getattr(FunctionsAIDecompilationApi, method)) +def test_line_attributions_carry_the_correspondence_map(): + assert "disassembly_line_number_to_ai_decompilation_line_numbers" in ( + LineAttributionsData.model_fields + ) + + +def test_line_attributions_accepts_the_kwargs_the_plugin_passes(): + params = inspect.signature( + FunctionsAIDecompilationApi.v3_get_ai_decompilation_line_attributions + ).parameters + assert "function_id" in params + + +def test_the_v3_blocks_method_is_the_one_the_plugin_calls(): + params = inspect.signature(FunctionsCoreApi.get_function_blocks_0).parameters + assert "function_id" in params + + serialize = inspect.getsource(FunctionsCoreApi._get_function_blocks_0_serialize) + assert "/v3/functions/{function_id}/blocks" in serialize + + +def test_blocks_response_exposes_basic_blocks(): + assert "basic_blocks" in DisassemblyOutputBody.model_fields + + def test_v3_token_methods_accept_plugin_kwargs(): tokens = inspect.signature( FunctionsAIDecompilationApi.v3_get_ai_decompilation_tokens diff --git a/tests/unit/ai_decomp/test_service.py b/tests/unit/ai_decomp/test_service.py index 63453f6..44c6b36 100644 --- a/tests/unit/ai_decomp/test_service.py +++ b/tests/unit/ai_decomp/test_service.py @@ -3,6 +3,7 @@ import pathlib import threading import time +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -428,6 +429,117 @@ def test_tokenised_phase_skipped_when_no_callback(service, sdk): sdk.v3_get_ai_decompilation_tokens.assert_not_called() +@pytest.fixture +def core_sdk(mocker): + api_class = mocker.patch.object(svc_mod, "FunctionsCoreApi") + api_inst = MagicMock() + api_class.return_value = api_inst + api_inst.get_function_blocks_0.return_value = SimpleNamespace( + basic_blocks=[{"min_addr": 0x1000, "asm": ["0x1000 push rbp", "0x1004 ret"]}] + ) + return api_inst + + +def _attributions(forward): + return SimpleNamespace( + disassembly_line_number_to_ai_decompilation_line_numbers=forward + ) + + +def _start_attribution(service, on_attributions): + service.start_ai_decomp_task( + ea=4096, + on_decomp=MagicMock(), + on_summary=MagicMock(), + on_comments=MagicMock(), + on_attributions=on_attributions, + ) + _wait(service) + + +def test_attribution_phase_builds_and_caches_the_map(service, sdk, core_sdk): + sdk.get_ai_decompilation.return_value = _dd(code="ok") + sdk.v3_get_ai_decompilation_line_attributions.return_value = _attributions({"0": [2]}) + + on_attributions = MagicMock() + _start_attribution(service, on_attributions) + + on_attributions.assert_called_once() + mapping = on_attributions.call_args[0][0] + assert mapping.addresses_for_decomp_line(2) == [0x1004] + assert mapping.decomp_lines_for_address(0x1004) == [2] + assert service._attribution_cache[42] is mapping + sdk.v3_get_ai_decompilation_line_attributions.assert_called_once_with(function_id=42) + core_sdk.get_function_blocks_0.assert_called_once_with(function_id=42) + + +def test_attribution_phase_serves_the_cache_without_refetching(service, sdk, core_sdk): + sdk.get_ai_decompilation.return_value = _dd(code="ok") + sdk.v3_get_ai_decompilation_line_attributions.return_value = _attributions({"0": [2]}) + + _start_attribution(service, MagicMock()) + second = MagicMock() + _start_attribution(service, second) + + second.assert_called_once() + assert sdk.v3_get_ai_decompilation_line_attributions.call_count == 1 + assert core_sdk.get_function_blocks_0.call_count == 1 + + +def test_an_empty_correspondence_skips_the_blocks_request(service, sdk, core_sdk): + sdk.get_ai_decompilation.return_value = _dd(code="ok") + sdk.v3_get_ai_decompilation_line_attributions.return_value = _attributions({}) + + on_attributions = MagicMock() + _start_attribution(service, on_attributions) + + on_attributions.assert_not_called() + core_sdk.get_function_blocks_0.assert_not_called() + assert service._attribution_cache[42].is_empty() + + +def test_a_failed_attribution_fetch_is_quiet(service, sdk, core_sdk): + sdk.get_ai_decompilation.return_value = _dd(code="ok") + sdk.v3_get_ai_decompilation_line_attributions.side_effect = ApiException(status=404) + + on_attributions = MagicMock() + _start_attribution(service, on_attributions) + + on_attributions.assert_not_called() + core_sdk.get_function_blocks_0.assert_not_called() + assert 42 not in service._attribution_cache + + +def test_a_failed_blocks_fetch_leaves_no_half_built_map(service, sdk, core_sdk): + sdk.get_ai_decompilation.return_value = _dd(code="ok") + sdk.v3_get_ai_decompilation_line_attributions.return_value = _attributions({"0": [2]}) + core_sdk.get_function_blocks_0.side_effect = ApiException(status=500) + + on_attributions = MagicMock() + _start_attribution(service, on_attributions) + + on_attributions.assert_not_called() + assert 42 not in service._attribution_cache + + +def test_attribution_phase_skipped_when_no_callback(service, sdk, core_sdk): + sdk.get_ai_decompilation.return_value = _dd(code="ok") + _run(service) + + sdk.v3_get_ai_decompilation_line_attributions.assert_not_called() + core_sdk.get_function_blocks_0.assert_not_called() + + +def test_invalidate_drops_the_cached_attribution_map(service, sdk, core_sdk): + sdk.get_ai_decompilation.return_value = _dd(code="ok") + sdk.v3_get_ai_decompilation_line_attributions.return_value = _attributions({"0": [2]}) + + _start_attribution(service, MagicMock()) + service.invalidate_ea(4096) + + assert 42 not in service._attribution_cache + + def test_apply_overrides_sends_body_refetches_and_caches(service, sdk): sdk.v3_upsert_ai_decompilation_overrides.return_value = MagicMock() sdk.get_ai_decompilation.return_value = _dd(code="renamed") @@ -741,7 +853,7 @@ def test_streaming_never_dispatches_on_the_main_thread(service, sdk, monkeypatch assert all(t is not threading.main_thread() for t in threads) -@pytest.mark.parametrize("callback", ["_on_stream", "_on_progress"]) +@pytest.mark.parametrize("callback", ["_on_stream", "_on_progress", "_on_attributions"]) def test_view_writing_callbacks_stay_marshalled_onto_the_ui_thread(callback): source = pathlib.Path( "reai_toolkit/app/coordinators/ai_decomp_coordinator.py" From 98e59e31662b11c7695c3b0fbd5c7000abfc6f6e Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:00:17 +0100 Subject: [PATCH 09/12] feat: vendor deps as script --- .github/workflows/publish.yaml | 5 +---- scripts/build_vendor.sh | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 4 deletions(-) create mode 100755 scripts/build_vendor.sh diff --git a/.github/workflows/publish.yaml b/.github/workflows/publish.yaml index 07099fe..4b219d8 100644 --- a/.github/workflows/publish.yaml +++ b/.github/workflows/publish.yaml @@ -67,10 +67,7 @@ jobs: - name: Build vendor dependencies run: | rm -rf screenshots - mkdir -p reai_toolkit/vendor - uv pip install . --target reai_toolkit/vendor --no-cache-dir - rm -rf reai_toolkit/vendor/reai_toolkit - rm -rf reai_toolkit/vendor/reai_toolkit-*.dist-info + bash scripts/build_vendor.sh shell: bash - name: Package plugin for release (cross-platform) diff --git a/scripts/build_vendor.sh b/scripts/build_vendor.sh new file mode 100755 index 0000000..56fbe22 --- /dev/null +++ b/scripts/build_vendor.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +set -euo pipefail + +cd "$(dirname "${BASH_SOURCE[0]}")/.." + +VENDOR="reai_toolkit/vendor" + +if ! command -v uv >/dev/null 2>&1; then + echo "uv is not on PATH; see https://docs.astral.sh/uv/getting-started/install/" >&2 + exit 1 +fi + +echo "==> Rebuilding vendor tree at $VENDOR" +rm -rf "$VENDOR" +mkdir -p "$VENDOR" +uv pip install . --target "$VENDOR" --no-cache-dir + +echo "==> Dropping the self-install so the checkout is not shadowed" +rm -rf "$VENDOR/reai_toolkit" "$VENDOR"/reai_toolkit-*.dist-info + +echo "==> Vendored (this is what the plugin and the tests import)" +for dist in "$VENDOR"/*.dist-info; do + [ -d "$dist" ] || continue + name="$(basename "$dist" .dist-info)" + case "$name" in + revengai-* | libbs-* | pydantic-* | loguru-* | requests-*) + echo " ${name%-*} ${name##*-}" + ;; + esac +done + +echo "==> $VENDOR is prepended to sys.path ahead of site-packages," +echo " so rerun this after any dependency change or the old tree wins." From 17996da50738545a683deef5e2d5a3c8027fe0bf Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:25:27 +0100 Subject: [PATCH 10/12] fix: test updates + gui bugs --- .../app/components/tabs/ai_decomp_tab.py | 60 +++- tests/ui/ai_decomp_ui_runner.py | 70 ++++- tests/ui/chat_ui_runner.py | 261 ++++++++++++++++++ tests/ui/test_gui_ai_decomp.py | 4 + tests/ui/test_gui_chat.py | 111 ++++++++ tests/unit/chat/test_qt_signal_receivers.py | 78 ++++++ 6 files changed, 570 insertions(+), 14 deletions(-) create mode 100644 tests/ui/chat_ui_runner.py create mode 100644 tests/ui/test_gui_chat.py create mode 100644 tests/unit/chat/test_qt_signal_receivers.py diff --git a/reai_toolkit/app/components/tabs/ai_decomp_tab.py b/reai_toolkit/app/components/tabs/ai_decomp_tab.py index 7938678..af4ffba 100644 --- a/reai_toolkit/app/components/tabs/ai_decomp_tab.py +++ b/reai_toolkit/app/components/tabs/ai_decomp_tab.py @@ -6,7 +6,7 @@ import ida_kernwin as kw from libbs.decompilers.ida.compat import execute_ui -from reai_toolkit.app.core.qt_compat import QtCore, QtGui, QtWidgets, Signal +from reai_toolkit.app.core.qt_compat import QtCore, QtGui, QtWidgets, Signal, Slot _WORD_UNDER_CURSOR = getattr( @@ -77,6 +77,44 @@ def contextMenuEvent(self, event) -> None: self.commentRemoveRequested.emit(line) +class _ViewBridge(QtCore.QObject): + def __init__(self, view: "AIDecompView") -> None: + super().__init__() + self._view = view + + @Slot() + def refresh(self) -> None: + self._view._on_refresh_clicked() + + @Slot() + def rate_up(self) -> None: + self._view._on_rate_up_clicked() + + @Slot() + def rate_down(self) -> None: + self._view._on_rate_down_clicked() + + @Slot() + def use_predicted_name(self) -> None: + self._view._on_use_predicted_name_clicked() + + @Slot(int, str) + def rename(self, line: int, word: str) -> None: + self._view._on_rename_requested(line, word) + + @Slot(int) + def edit_comment(self, line: int) -> None: + self._view._on_edit_comment_requested(line) + + @Slot(int) + def remove_comment(self, line: int) -> None: + self._view._on_remove_comment_requested(line) + + @Slot() + def cursor_moved(self) -> None: + self._view._on_cursor_moved() + + class AIDecompView(kw.PluginForm): """ Dockable tab using Qt editor + QSyntaxHighlighter. @@ -108,6 +146,7 @@ def __init__(self, on_closed: Optional[Callable[[], None]] = None) -> None: self._predicted_label: QtWidgets.QLabel | None = None self._predicted_btn: QtWidgets.QPushButton | None = None self._predicted_name: str | None = None + self._bridge: _ViewBridge | None = None self._highlighter: CppHighlighter | None = None def Create(self, title: Any) -> Any: @@ -132,6 +171,8 @@ def OnCreate(self, form) -> None: """Called by IDA when the form is created; build our Qt UI here.""" self._parent_window = self.FormToPyQtWidget(form) + self._bridge = _ViewBridge(self) + # Layout root layout = QtWidgets.QVBoxLayout(self._parent_window) layout.setContentsMargins(0, 0, 0, 0) @@ -150,7 +191,7 @@ def OnCreate(self, form) -> None: ) self._predicted_btn.setToolTip("Rename this function to the predicted name") self._predicted_btn.setVisible(False) - self._predicted_btn.clicked.connect(self._on_use_predicted_name_clicked) + self._predicted_btn.clicked.connect(self._bridge.use_predicted_name) header.addWidget(self._predicted_btn) header.addStretch(1) @@ -164,7 +205,7 @@ def OnCreate(self, form) -> None: self._rate_up_btn.setText("\U0001f44d") self._rate_up_btn.setCheckable(True) self._rate_up_btn.setToolTip("Rate this AI decompilation as good") - self._rate_up_btn.clicked.connect(self._on_rate_up_clicked) + self._rate_up_btn.clicked.connect(self._bridge.rate_up) header.addWidget(self._rate_up_btn) down_icon = _thumb_icon(up=False) @@ -176,11 +217,11 @@ def OnCreate(self, form) -> None: self._rate_down_btn.setText("\U0001f44e") self._rate_down_btn.setCheckable(True) self._rate_down_btn.setToolTip("Rate this AI decompilation as poor") - self._rate_down_btn.clicked.connect(self._on_rate_down_clicked) + self._rate_down_btn.clicked.connect(self._bridge.rate_down) header.addWidget(self._rate_down_btn) self._refresh_btn = QtWidgets.QPushButton("Refresh", self._parent_window) - self._refresh_btn.clicked.connect(self._on_refresh_clicked) + self._refresh_btn.clicked.connect(self._bridge.refresh) header.addWidget(self._refresh_btn) layout.addLayout(header) @@ -188,10 +229,10 @@ def OnCreate(self, form) -> None: self._editor = _DecompEditor(self._parent_window) self._editor.setReadOnly(True) self._editor.setLineWrapMode(QtWidgets.QPlainTextEdit.NoWrap) - self._editor.renameRequested.connect(self._on_rename_requested) - self._editor.commentEditRequested.connect(self._on_edit_comment_requested) - self._editor.commentRemoveRequested.connect(self._on_remove_comment_requested) - self._editor.cursorPositionChanged.connect(self._on_cursor_moved) + self._editor.renameRequested.connect(self._bridge.rename) + self._editor.commentEditRequested.connect(self._bridge.edit_comment) + self._editor.commentRemoveRequested.connect(self._bridge.remove_comment) + self._editor.cursorPositionChanged.connect(self._bridge.cursor_moved) # Monospace font tuned for IDA font = QtGui.QFont( @@ -225,6 +266,7 @@ def OnClose(self, form) -> None: self._predicted_label = None self._predicted_btn = None self._predicted_name = None + self._bridge = None self._parent_window = None def _on_refresh_clicked(self) -> None: diff --git a/tests/ui/ai_decomp_ui_runner.py b/tests/ui/ai_decomp_ui_runner.py index 6810eca..b28ad0b 100644 --- a/tests/ui/ai_decomp_ui_runner.py +++ b/tests/ui/ai_decomp_ui_runner.py @@ -32,6 +32,7 @@ def _run(report: dict) -> None: from unittest.mock import MagicMock import ida_kernwin + import idaapi import idautils from revengai.models.comments_data import CommentsData @@ -279,11 +280,10 @@ def code_line_row(needle: str) -> int: coord.on_disassembly_ea(ea) pump() - report["disassembly_ea_lights_its_decomp_line"] = ( - len(view._editor.extraSelections()) == 1 - and view._editor.extraSelections()[0].cursor.blockNumber() - == code_line_row("int v5") - ) + selections = view._editor.extraSelections() + report["disassembly_ea_lights_its_decomp_line"] = [ + sel.cursor.blockNumber() for sel in selections + ] == [code_line_row("int v5")] coord.on_disassembly_ea(0xDEAD) pump() @@ -294,6 +294,62 @@ def code_line_row(needle: str) -> int: coord._attribution_hooks = None coord._current_attributions = None + from reai_toolkit.hooks.reactive import LineAttributionHooks + + class _CountingHooks(LineAttributionHooks): + def __init__(self, coordinator): + super().__init__(coordinator) + self.calls = 0 + self.painted = 0 + self.offered = [] + self.widget_types = [] + + def reset(self): + self.calls = 0 + self.painted = 0 + self.offered = [] + + def get_lines_rendering_info(self, out, widget, rin): + self.calls += 1 + self.widget_types.append(ida_kernwin.get_widget_type(widget)) + for section in rin.sections_lines: + for line in section: + self.offered.append(line.at.toea()) + before = out.entries.size() + super().get_lines_rendering_info(out, widget, rin) + self.painted += out.entries.size() - before + + def repaint(target=None): + if target is not None: + ida_kernwin.jumpto(target) + ida_kernwin.refresh_idaview_anyway() + pump() + + wanted = [] + for func in list(idautils.Functions())[:40]: + wanted.extend(idautils.FuncItems(func)) + wanted_set = set(wanted) + + probe = _CountingHooks(coord) + probe.set_addresses(wanted) + probe.hook() + repaint(next(iter(idautils.Functions()), ea)) + + offered = [a for a in probe.offered if a != idaapi.BADADDR] + expected = sum(1 for a in offered if a in wanted_set) + + report["diag_offered"] = len(offered) + report["diag_expected"] = expected + report["diag_painted"] = probe.painted + + report["rendering_hook_runs_on_repaint"] = probe.calls > 0 + report["rendering_hook_sees_the_disassembly"] = bool(probe.widget_types) and set( + probe.widget_types + ) == {ida_kernwin.BWN_DISASM} + report["attributed_addresses_paint"] = expected > 0 and probe.painted == expected + report["unattributed_lines_are_left_alone"] = probe.painted < len(offered) + probe.unhook() + service.reset_mock() service.peek_decomp.return_value = None view._refresh_btn.click() @@ -333,6 +389,10 @@ def main() -> None: "a_line_with_no_counterpart_lights_nothing": False, "disassembly_ea_lights_its_decomp_line": False, "an_unattributed_address_lights_nothing": False, + "rendering_hook_runs_on_repaint": False, + "rendering_hook_sees_the_disassembly": False, + "attributed_addresses_paint": False, + "unattributed_lines_are_left_alone": False, "refresh_button_invalidates": False, } ida_auto.auto_wait() diff --git a/tests/ui/chat_ui_runner.py b/tests/ui/chat_ui_runner.py new file mode 100644 index 0000000..d70a0b2 --- /dev/null +++ b/tests/ui/chat_ui_runner.py @@ -0,0 +1,261 @@ +import json +import os +import sys +import threading +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: + import time + + from reai_toolkit.app.components.tabs.chat_tab import ChatPanel, ChatStreamWorker + from reai_toolkit.app.core.qt_compat import QtCore, QtGui, QtWidgets + from reai_toolkit.app.core.shared_schema import GenericApiReturn + from reai_toolkit.app.services.chat.reducer import ( + EventAction, + SendMessage, + chat_reducer, + initial_state, + ) + from reai_toolkit.app.services.chat.schema import ChatEvent + + qapp = QtWidgets.QApplication.instance() + + def pump(rounds: int = 10) -> None: + if qapp is not None: + for _ in range(rounds): + qapp.processEvents() + + def pump_until(predicate, timeout: float = 10.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + pump() + if predicate(): + return True + time.sleep(0.01) + return predicate() + + panel = ChatPanel(on_closed=lambda: None) + panel.Create(ChatPanel.TITLE) + pump() + + if panel._transcript is None or panel._input is None: + report["errors"].append("ChatPanel widgets were not created") + return + + report["panel_created"] = True + report["bridge_is_a_qobject"] = isinstance(panel._bridge, QtCore.QObject) + report["relay_is_a_qobject"] = isinstance(panel._relay, QtCore.QObject) + report["relay_lives_on_the_ui_thread"] = ( + panel._relay.thread() is QtCore.QThread.currentThread() + ) + + fired: list = [] + panel.on_send = lambda text: fired.append(("send", text)) + panel.on_stop = lambda: fired.append(("stop", None)) + panel.on_new_chat = lambda: fired.append(("new_chat", None)) + panel.on_request_history = lambda: fired.append(("history", None)) + panel.on_confirm = lambda cid, ok: fired.append(("confirm", (cid, ok))) + panel.on_select_conversation = lambda uuid: fired.append(("select", uuid)) + panel.on_jump = lambda ea: fired.append(("jump", ea)) + + panel._input.setPlainText("what does this do?") + panel._send_btn.click() + pump() + report["send_button_routes_through_the_bridge"] = fired[-1:] == [ + ("send", "what does this do?") + ] + report["send_clears_the_input"] = panel._input.toPlainText() == "" + + fired.clear() + panel._stop_btn.setEnabled(True) + panel._stop_btn.click() + pump() + report["stop_button_routes_through_the_bridge"] = fired[-1:] == [("stop", None)] + + fired.clear() + panel._input.setPlainText(" ") + panel._send_btn.click() + pump() + report["blank_input_sends_nothing"] = fired == [] + panel._input.clear() + + fired.clear() + panel._input.setPlainText("typed") + panel._input.keyPressEvent( + QtGui.QKeyEvent( + QtCore.QEvent.KeyPress, QtCore.Qt.Key_Return, QtCore.Qt.NoModifier + ) + ) + pump() + report["enter_submits_through_the_bridge"] = fired[-1:] == [("send", "typed")] + + state = initial_state() + state = chat_reducer(state, SendMessage(id="u1", content="hello there")) + state = chat_reducer( + state, EventAction(ChatEvent(type="TITLE_UPDATED", title="A Conversation")) + ) + panel.request_render(state) + report["render_is_deferred_to_the_timer"] = panel._render_timer.isActive() + pump_until(lambda: not panel._render_timer.isActive()) + report["timer_flush_reaches_the_transcript"] = ( + "hello there" in panel._transcript.toPlainText() + ) + report["timer_flush_updates_the_title"] = panel._title_label.text() == "A Conversation" + + confirming = chat_reducer( + state, + EventAction( + ChatEvent( + type="TOOL_CONFIRMATION_REQUIRED", + tool_call_id="tc1", + tool_name="rename_functions", + message="Rename 3 functions?", + ) + ), + ) + panel.request_render(confirming) + pump_until(lambda: not panel._render_timer.isActive()) + report["confirmation_shows_its_bar"] = ( + not panel._confirm_bar.isHidden() and "Rename 3" in panel._confirm_msg.text() + ) + + fired.clear() + panel._confirm_bar.findChildren(QtWidgets.QPushButton)[0].click() + pump() + report["approve_routes_with_the_pending_id"] = fired[-1:] == [ + ("confirm", ("tc1", True)) + ] + + events = [ + ChatEvent(type="TEXT_MESSAGE_START", message_id="m1", role="assistant"), + ChatEvent(type="TEXT_MESSAGE_CONTENT", message_id="m1", delta="streamed"), + ChatEvent(type="TEXT_MESSAGE_END", message_id="m1"), + ChatEvent(type="RUN_FINISHED"), + ] + + class _FakeService: + def __init__(self): + self.closed = False + + def create_conversation(self, context): + return GenericApiReturn(success=True, data="conv-uuid") + + def send_message(self, conv_id, content, context): + return GenericApiReturn(success=True) + + def stream(self, conv_id, stop_event, last_event_id=None): + for event in events: + if stop_event.is_set(): + return + yield event + + def close_active_stream(self): + self.closed = True + + seen: list = [] + threads: list = [] + created: list = [] + finished: list = [] + panel.on_stream_event = lambda ev: ( + seen.append(ev.type), + threads.append(threading.current_thread()), + ) + panel.on_stream_conversation_created = lambda uuid: created.append(uuid) + panel.on_stream_finished = lambda: finished.append(True) + panel.on_stream_error = lambda msg: report["errors"].append(f"stream errored: {msg}") + + worker = ChatStreamWorker(_FakeService(), None, "hello", None) + panel.start_stream_worker(worker) + report["worker_reports_streaming"] = panel.is_streaming() + + pump_until(lambda: bool(finished)) + + report["every_streamed_event_arrives"] = seen == [e.type for e in events] + report["stream_callbacks_run_on_the_ui_thread"] = bool(threads) and all( + t is threading.main_thread() for t in threads + ) + report["conversation_id_reaches_the_panel"] = created == ["conv-uuid"] + report["stream_finish_clears_streaming"] = not panel.is_streaming() + + class _BlockingService(_FakeService): + def stream(self, conv_id, stop_event, last_event_id=None): + while not stop_event.is_set(): + time.sleep(0.01) + yield from () + + blocking = _BlockingService() + panel.start_stream_worker(ChatStreamWorker(blocking, "conv", None, None)) + pump() + started = time.monotonic() + panel.stop_stream_worker() + report["stop_returns_promptly"] = (time.monotonic() - started) < 3.0 + report["stop_closes_the_active_stream"] = blocking.closed + report["stop_leaves_no_zombie_thread"] = panel._zombies == [] + + panel.Close(0) + pump() + report["ok"] = not report["errors"] + + +def main() -> None: + import ida_auto + import ida_pro + + report = { + "ok": False, + "errors": [], + "panel_created": False, + "bridge_is_a_qobject": False, + "relay_is_a_qobject": False, + "relay_lives_on_the_ui_thread": False, + "send_button_routes_through_the_bridge": False, + "send_clears_the_input": False, + "stop_button_routes_through_the_bridge": False, + "blank_input_sends_nothing": False, + "enter_submits_through_the_bridge": False, + "render_is_deferred_to_the_timer": False, + "timer_flush_reaches_the_transcript": False, + "timer_flush_updates_the_title": False, + "confirmation_shows_its_bar": False, + "approve_routes_with_the_pending_id": False, + "worker_reports_streaming": False, + "every_streamed_event_arrives": False, + "stream_callbacks_run_on_the_ui_thread": False, + "conversation_id_reaches_the_panel": False, + "stream_finish_clears_streaming": False, + "stop_returns_promptly": False, + "stop_closes_the_active_stream": False, + "stop_leaves_no_zombie_thread": False, + } + 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_ai_decomp.py b/tests/ui/test_gui_ai_decomp.py index 6a22505..3ef225f 100644 --- a/tests/ui/test_gui_ai_decomp.py +++ b/tests/ui/test_gui_ai_decomp.py @@ -36,6 +36,10 @@ "a_line_with_no_counterpart_lights_nothing", "disassembly_ea_lights_its_decomp_line", "an_unattributed_address_lights_nothing", + "rendering_hook_runs_on_repaint", + "rendering_hook_sees_the_disassembly", + "attributed_addresses_paint", + "unattributed_lines_are_left_alone", "refresh_button_invalidates", ] diff --git a/tests/ui/test_gui_chat.py b/tests/ui/test_gui_chat.py new file mode 100644 index 0000000..b5f29f1 --- /dev/null +++ b/tests/ui/test_gui_chat.py @@ -0,0 +1,111 @@ +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__)), "chat_ui_runner.py") +HELLO_ELF = os.path.join(ROOT, "tests", "fixtures", "hello.elf") + +CHECKS = [ + "panel_created", + "bridge_is_a_qobject", + "relay_is_a_qobject", + "relay_lives_on_the_ui_thread", + "send_button_routes_through_the_bridge", + "send_clears_the_input", + "stop_button_routes_through_the_bridge", + "blank_input_sends_nothing", + "enter_submits_through_the_bridge", + "render_is_deferred_to_the_timer", + "timer_flush_reaches_the_transcript", + "timer_flush_updates_the_title", + "confirmation_shows_its_bar", + "approve_routes_with_the_pending_id", + "worker_reports_streaming", + "every_streamed_event_arrives", + "stream_callbacks_run_on_the_ui_thread", + "conversation_id_reaches_the_panel", + "stream_finish_clears_streaming", + "stop_returns_promptly", + "stop_closes_the_active_stream", + "stop_leaves_no_zombie_thread", +] + + +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_agent_chat_panel_in_gui(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"] == [], report["errors"] + assert report["ok"] is True + for check in CHECKS: + assert report.get(check) is True, f"{check} failed: {report}" diff --git a/tests/unit/chat/test_qt_signal_receivers.py b/tests/unit/chat/test_qt_signal_receivers.py new file mode 100644 index 0000000..74fe75f --- /dev/null +++ b/tests/unit/chat/test_qt_signal_receivers.py @@ -0,0 +1,78 @@ +"""Guards the PySide6 6.8 abort: signals must target a QObject, never a PluginForm.""" + +import ast +import pathlib + +import pytest + +PANELS = [ + "reai_toolkit/app/components/tabs/chat_tab.py", + "reai_toolkit/app/components/tabs/ai_decomp_tab.py", +] + +QOBJECT_RECEIVERS = ("_bridge", "_relay") + + +def _connect_calls(tree): + for node in ast.walk(tree): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "connect" + and node.args + ): + yield node + + +def _describe(arg) -> str: + return ast.unparse(arg) + + +def _is_qobject_receiver(arg) -> bool: + if not isinstance(arg, ast.Attribute): + return False + base = arg.value + if isinstance(base, ast.Name): + return base.id != "self" + if ( + isinstance(base, ast.Attribute) + and isinstance(base.value, ast.Name) + and base.value.id == "self" + ): + return base.attr in QOBJECT_RECEIVERS + return False + + +@pytest.mark.parametrize("path", PANELS) +def test_panel_signals_are_delivered_to_a_qobject(path): + source = pathlib.Path(path).read_text() + tree = ast.parse(source) + + offenders = [ + f"line {call.lineno}: connect({_describe(call.args[0])})" + for call in _connect_calls(tree) + if not _is_qobject_receiver(call.args[0]) + ] + + assert offenders == [], ( + "PySide6 6.8 aborts building a global receiver when a signal is wired to a " + "callable on a non-QObject PluginForm; route it through a @Slot on a QObject " + f"bridge instead ({path}): {offenders}" + ) + + +@pytest.mark.parametrize("path", PANELS) +def test_each_panel_actually_owns_a_qobject_bridge(path): + tree = ast.parse(pathlib.Path(path).read_text()) + + bridges = [ + node.name + for node in ast.walk(tree) + if isinstance(node, ast.ClassDef) + and any( + isinstance(base, ast.Attribute) and base.attr == "QObject" + for base in node.bases + ) + ] + + assert bridges, f"{path} connects signals but defines no QObject receiver" From ca9e4d6f2d3234e3b5431db8fee62cfd15704c62 Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:46:45 +0100 Subject: [PATCH 11/12] fix: auth bug --- .../app/services/auth/auth_service.py | 2 + tests/unit/auth/test_endpoint_auth.py | 119 ++++++++++++++++++ tests/unit/auth/test_service.py | 23 ++++ 3 files changed, 144 insertions(+) create mode 100644 tests/unit/auth/test_endpoint_auth.py diff --git a/reai_toolkit/app/services/auth/auth_service.py b/reai_toolkit/app/services/auth/auth_service.py index 51ee81b..071b5ab 100644 --- a/reai_toolkit/app/services/auth/auth_service.py +++ b/reai_toolkit/app/services/auth/auth_service.py @@ -55,11 +55,13 @@ def build_sdk_config(self) -> Configuration: self.sdk_config = Configuration( host=self.config_service.api_url, api_key={"APIKey": self.config_service.api_key}, + access_token=self.config_service.api_key, ) # Other calls must reassign values - do not create new object (mutate existing one, all holders share ref) else: self.sdk_config.host = self.config_service.api_url self.sdk_config.api_key = {"APIKey": self.config_service.api_key} + self.sdk_config.access_token = self.config_service.api_key self.sdk_config.user_agent = ( f"IDA/{self._ida_version} RevEng.AI_Plugin/{self._plugin_version}" diff --git a/tests/unit/auth/test_endpoint_auth.py b/tests/unit/auth/test_endpoint_auth.py new file mode 100644 index 0000000..a12294e --- /dev/null +++ b/tests/unit/auth/test_endpoint_auth.py @@ -0,0 +1,119 @@ +"""Every SDK endpoint the plugin calls must be covered by a configured auth scheme. + +A request whose declared scheme the plugin never configures goes out with no +Authorization header at all and the server answers 401 — which is what happened +to the bearerAuth-only /v2/iam/me. +""" + +import ast +import collections +import inspect +import pathlib +import re + +import pytest +import revengai +from unittest.mock import MagicMock + +from reai_toolkit.app.services.auth.auth_service import AuthService + +PLUGIN_ROOT = pathlib.Path("reai_toolkit/app") + + +@pytest.fixture +def service(): + cfg = MagicMock() + cfg.api_url = "https://api.reveng.ai" + cfg.api_key = "secret-key" + return AuthService(cfg=cfg, ida_version="9.3", plugin_version="1.2.3") + + +_AUTH_RE = re.compile(r"_auth_settings:\s*List\[str\]\s*=\s*\[([^\]]*)\]", re.S) +_VARIANT_SUFFIXES = ("_without_preload_content", "_with_http_info") + + +def _api_calls() -> dict[str, set[str]]: + calls: dict[str, set[str]] = collections.defaultdict(set) + for path in PLUGIN_ROOT.rglob("*.py"): + tree = ast.parse(path.read_text()) + aliases: dict[str, str] = {} + for node in ast.walk(tree): + if isinstance(node, ast.Assign) and isinstance(node.value, ast.Call): + func = node.value.func + if isinstance(func, ast.Name) and func.id.endswith("Api"): + for target in node.targets: + if isinstance(target, ast.Name): + aliases[target.id] = func.id + for node in ast.walk(tree): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)): + continue + receiver = node.func.value + api = None + if ( + isinstance(receiver, ast.Call) + and isinstance(receiver.func, ast.Name) + and receiver.func.id.endswith("Api") + ): + api = receiver.func.id + elif isinstance(receiver, ast.Name) and receiver.id in aliases: + api = aliases[receiver.id] + if api is not None: + calls[api].add(node.func.attr) + return calls + + +def _declared_schemes(api_name: str, method: str) -> set[str] | None: + api = getattr(revengai, api_name, None) + if api is None: + return None + for suffix in _VARIANT_SUFFIXES: + if method.endswith(suffix): + method = method[: -len(suffix)] + break + serialize = getattr(api, f"_{method}_serialize", None) + if serialize is None: + return None + match = _AUTH_RE.search(inspect.getsource(serialize)) + if match is None: + return None + return {s.strip().strip("'\"") for s in match.group(1).split(",") if s.strip()} + + +def test_the_call_discovery_actually_finds_the_plugins_endpoints(): + calls = _api_calls() + + assert "get_me" in calls.get("IAMUsersApi", set()) + assert "get_config" in calls.get("ConfigApi", set()) + assert sum(len(v) for v in calls.values()) > 40 + + +def test_every_endpoint_the_plugin_calls_carries_credentials(service): + configured = set(service.build_sdk_config().auth_settings()) + + unauthenticated = [ + f"{api}.{method} declares {sorted(declared)}" + for api, methods in sorted(_api_calls().items()) + for method in sorted(methods) + if (declared := _declared_schemes(api, method)) and not declared & configured + ] + + assert unauthenticated == [], ( + f"the plugin configures {sorted(configured)}; these endpoints declare none " + f"of them, so their requests are sent unauthenticated: {unauthenticated}" + ) + + +def test_bearer_auth_is_load_bearing(service): + without_bearer = set(service.build_sdk_config().auth_settings()) - {"bearerAuth"} + + stranded = [ + f"{api}.{method}" + for api, methods in _api_calls().items() + for method in methods + if (declared := _declared_schemes(api, method)) and not declared & without_bearer + ] + + assert stranded, ( + "nothing requires bearerAuth any more, so the coverage test above proves " + "nothing; either access_token can be dropped or the discovery has broken" + ) diff --git a/tests/unit/auth/test_service.py b/tests/unit/auth/test_service.py index 300e2d4..886f1fc 100644 --- a/tests/unit/auth/test_service.py +++ b/tests/unit/auth/test_service.py @@ -49,6 +49,7 @@ def test_build_sdk_config_sets_host_key_and_user_agent(service): assert cfg.host == "https://api.reveng.ai" assert cfg.api_key == {"APIKey": "secret-key"} + assert cfg.access_token == "secret-key" assert cfg.user_agent == "IDA/9.3 RevEng.AI_Plugin/1.2.3" @@ -62,6 +63,7 @@ def test_get_sdk_config_builds_once_and_mutates_in_place(service): assert first is second assert second.host == "https://other.host" assert second.api_key == {"APIKey": "new-key"} + assert second.access_token == "new-key" def test_verify_success(service, config_api): @@ -169,3 +171,24 @@ def test_verify_success_warms_user_cache(service, config_api, iam_api): assert ok is True iam_api.get_me.assert_called_once() assert service.is_enthusiast() is True + + +def test_both_auth_schemes_are_configured(service): + schemes = service.build_sdk_config().auth_settings() + + assert set(schemes) == {"APIKey", "bearerAuth"} + + +def test_bearer_only_endpoints_still_receive_an_authorization_header(service): + from revengai.api_client import ApiClient + + client = ApiClient(configuration=service.build_sdk_config()) + headers: dict = {} + client.update_params_for_auth( + headers, [], ["bearerAuth"], "/v2/iam/me", "GET", None + ) + + assert headers["Authorization"] == "Bearer secret-key", ( + "/v2/iam/me accepts bearerAuth only; without access_token the request " + "goes out unauthenticated and the server answers 401" + ) From d11a52c9b9a719b735c101d3703b0fdfba11bcbc Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Wed, 16 Sep 2026 16:32:31 +0100 Subject: [PATCH 12/12] fix: testing patches --- .../app/components/dialogs/analyse_dialog.py | 1 + .../forms/analyse/analyse_panel_qt5.ui | 29 ++- .../forms/analyse/analyse_panel_ui_uic5.py | 27 ++- .../forms/analyse/analyse_panel_ui_uic6.py | 37 +++- .../app/components/tabs/chat_render.py | 55 ++++- reai_toolkit/app/components/tabs/chat_tab.py | 16 +- .../app/coordinators/ai_decomp_coordinator.py | 20 +- .../app/coordinators/chat_coordinator.py | 28 +++ reai_toolkit/app/interfaces/base_service.py | 5 + .../services/ai_decomp/ai_decomp_service.py | 1 - .../analysis_status/analysis_status.py | 21 +- .../services/analysis_sync/analysis_sync.py | 14 +- .../services/data_types/analysis_catalogue.py | 113 +++++++--- .../app/services/rename/rename_service.py | 80 +++++-- .../app/services/upload/upload_service.py | 9 +- .../unit/analysis_status/test_sdk_schemas.py | 15 +- tests/unit/analysis_status/test_service.py | 8 +- tests/unit/analysis_sync/test_sdk_schemas.py | 8 +- tests/unit/analysis_sync/test_service.py | 58 +++++ tests/unit/chat/test_render.py | 63 ++++++ tests/unit/core/test_sdk_call_signatures.py | 205 ++++++++++++++++++ .../data_types/test_analysis_catalogue.py | 22 +- tests/unit/rename/test_sdk_schemas.py | 18 +- tests/unit/rename/test_service.py | 36 ++- tests/unit/upload/test_service.py | 24 ++ 25 files changed, 776 insertions(+), 137 deletions(-) create mode 100644 tests/unit/core/test_sdk_call_signatures.py diff --git a/reai_toolkit/app/components/dialogs/analyse_dialog.py b/reai_toolkit/app/components/dialogs/analyse_dialog.py index 82c66d4..dc66c6c 100644 --- a/reai_toolkit/app/components/dialogs/analyse_dialog.py +++ b/reai_toolkit/app/components/dialogs/analyse_dialog.py @@ -179,6 +179,7 @@ def _on_ok(self) -> None: debug_file_path=self.debug_file_path, tags=self._parse_tags(), public=self.ui.radioButton_2.isChecked(), + no_cache=self.ui.noCacheCheckBox.isChecked(), thread_callback=self.callback, ) diff --git a/reai_toolkit/app/components/forms/analyse/analyse_panel_qt5.ui b/reai_toolkit/app/components/forms/analyse/analyse_panel_qt5.ui index a653a50..3b36ed5 100644 --- a/reai_toolkit/app/components/forms/analyse/analyse_panel_qt5.ui +++ b/reai_toolkit/app/components/forms/analyse/analyse_panel_qt5.ui @@ -7,19 +7,19 @@ 0 0 520 - 420 + 460 520 - 420 + 460 520 - 420 + 460 @@ -132,13 +132,13 @@ 496 - 200 + 240 496 - 200 + 240 @@ -336,13 +336,30 @@ + + + Cache + + + + + + + Skip cached results and re-run every processing stage. + + + Disable cache + + + + Functions - + diff --git a/reai_toolkit/app/components/forms/analyse/analyse_panel_ui_uic5.py b/reai_toolkit/app/components/forms/analyse/analyse_panel_ui_uic5.py index 15f1ef7..56ae719 100644 --- a/reai_toolkit/app/components/forms/analyse/analyse_panel_ui_uic5.py +++ b/reai_toolkit/app/components/forms/analyse/analyse_panel_ui_uic5.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -# Form implementation generated from reading ui file 'analyse_panel_qt5.ui' +# Form implementation generated from reading ui file 'reai_toolkit/app/components/forms/analyse/analyse_panel_qt5.ui' # # Created by: PyQt5 UI code generator 5.15.11 # @@ -14,9 +14,9 @@ class Ui_AuthPanel(object): def setupUi(self, AuthPanel): AuthPanel.setObjectName("AuthPanel") - AuthPanel.resize(520, 420) - AuthPanel.setMinimumSize(QtCore.QSize(520, 420)) - AuthPanel.setMaximumSize(QtCore.QSize(520, 420)) + AuthPanel.resize(520, 460) + AuthPanel.setMinimumSize(QtCore.QSize(520, 460)) + AuthPanel.setMaximumSize(QtCore.QSize(520, 460)) AuthPanel.setAutoFillBackground(False) AuthPanel.setStyleSheet("\n" " QGroupBox{font-weight:600;margin-top:8px;}\n" @@ -47,7 +47,7 @@ def setupUi(self, AuthPanel): self.logoArea.setMaximumSize(QtCore.QSize(96, 96)) self.logoArea.setFrameShape(QtWidgets.QFrame.NoFrame) self.logoArea.setText("") - self.logoArea.setPixmap(QtGui.QPixmap("../../resources/reveng_ai_logo.jpg")) + self.logoArea.setPixmap(QtGui.QPixmap("reai_toolkit/app/components/forms/analyse/../../resources/reveng_ai_logo.jpg")) self.logoArea.setScaledContents(True) self.logoArea.setObjectName("logoArea") self.headerLayout.addWidget(self.logoArea) @@ -65,8 +65,8 @@ def setupUi(self, AuthPanel): self.headerLayout.addLayout(self.titleLayout) self.rootLayout.addWidget(self.header) self.groupEndpoints = QtWidgets.QGroupBox(AuthPanel) - self.groupEndpoints.setMinimumSize(QtCore.QSize(496, 200)) - self.groupEndpoints.setMaximumSize(QtCore.QSize(496, 200)) + self.groupEndpoints.setMinimumSize(QtCore.QSize(496, 240)) + self.groupEndpoints.setMaximumSize(QtCore.QSize(496, 240)) self.groupEndpoints.setObjectName("groupEndpoints") self.formLayout = QtWidgets.QFormLayout(self.groupEndpoints) self.formLayout.setLabelAlignment(QtCore.Qt.AlignRight) @@ -137,15 +137,21 @@ def setupUi(self, AuthPanel): self.radioButton.setObjectName("radioButton") self.horizontalLayout.addWidget(self.radioButton) self.formLayout.setLayout(4, QtWidgets.QFormLayout.FieldRole, self.horizontalLayout) + self.labelCache = QtWidgets.QLabel(self.groupEndpoints) + self.labelCache.setObjectName("labelCache") + self.formLayout.setWidget(5, QtWidgets.QFormLayout.LabelRole, self.labelCache) + self.noCacheCheckBox = QtWidgets.QCheckBox(self.groupEndpoints) + self.noCacheCheckBox.setObjectName("noCacheCheckBox") + self.formLayout.setWidget(5, QtWidgets.QFormLayout.FieldRole, self.noCacheCheckBox) self.functionsLabel = QtWidgets.QLabel(self.groupEndpoints) self.functionsLabel.setObjectName("functionsLabel") - self.formLayout.setWidget(5, QtWidgets.QFormLayout.LabelRole, self.functionsLabel) + self.formLayout.setWidget(6, QtWidgets.QFormLayout.LabelRole, self.functionsLabel) self.selectFuncs = QtWidgets.QPushButton(self.groupEndpoints) self.selectFuncs.setMinimumSize(QtCore.QSize(296, 32)) self.selectFuncs.setMaximumSize(QtCore.QSize(296, 32)) self.selectFuncs.setAutoDefault(False) self.selectFuncs.setObjectName("selectFuncs") - self.formLayout.setWidget(5, QtWidgets.QFormLayout.FieldRole, self.selectFuncs) + self.formLayout.setWidget(6, QtWidgets.QFormLayout.FieldRole, self.selectFuncs) self.rootLayout.addWidget(self.groupEndpoints) self.linkRow = QtWidgets.QHBoxLayout() self.linkRow.setObjectName("linkRow") @@ -189,6 +195,9 @@ def retranslateUi(self, AuthPanel): self.labelScope.setText(_translate("AuthPanel", "Scope")) self.radioButton_2.setText(_translate("AuthPanel", "Public")) self.radioButton.setText(_translate("AuthPanel", "Private")) + self.labelCache.setText(_translate("AuthPanel", "Cache")) + self.noCacheCheckBox.setToolTip(_translate("AuthPanel", "Skip cached results and re-run every processing stage.")) + self.noCacheCheckBox.setText(_translate("AuthPanel", "Disable cache")) self.functionsLabel.setText(_translate("AuthPanel", "Functions")) self.selectFuncs.setText(_translate("AuthPanel", "Select Functions to Upload")) self.cancelButton.setText(_translate("AuthPanel", "Cancel")) diff --git a/reai_toolkit/app/components/forms/analyse/analyse_panel_ui_uic6.py b/reai_toolkit/app/components/forms/analyse/analyse_panel_ui_uic6.py index 2cd53aa..abc6c95 100644 --- a/reai_toolkit/app/components/forms/analyse/analyse_panel_ui_uic6.py +++ b/reai_toolkit/app/components/forms/analyse/analyse_panel_ui_uic6.py @@ -15,18 +15,18 @@ QFont, QFontDatabase, QGradient, QIcon, QImage, QKeySequence, QLinearGradient, QPainter, QPalette, QPixmap, QRadialGradient, QTransform) -from PySide6.QtWidgets import (QApplication, QFormLayout, QFrame, QGroupBox, - QHBoxLayout, QLabel, QLineEdit, QPushButton, - QRadioButton, QSizePolicy, QSpacerItem, QVBoxLayout, - QWidget) +from PySide6.QtWidgets import (QApplication, QCheckBox, QFormLayout, QFrame, + QGroupBox, QHBoxLayout, QLabel, QLineEdit, + QPushButton, QRadioButton, QSizePolicy, QSpacerItem, + QVBoxLayout, QWidget) class Ui_AuthPanel(object): def setupUi(self, AuthPanel): if not AuthPanel.objectName(): AuthPanel.setObjectName(u"AuthPanel") - AuthPanel.resize(520, 420) - AuthPanel.setMinimumSize(QSize(520, 420)) - AuthPanel.setMaximumSize(QSize(520, 420)) + AuthPanel.resize(520, 460) + AuthPanel.setMinimumSize(QSize(520, 460)) + AuthPanel.setMaximumSize(QSize(520, 460)) AuthPanel.setAutoFillBackground(False) AuthPanel.setStyleSheet(u"\n" " QGroupBox{font-weight:600;margin-top:8px;}\n" @@ -85,8 +85,8 @@ def setupUi(self, AuthPanel): self.groupEndpoints = QGroupBox(AuthPanel) self.groupEndpoints.setObjectName(u"groupEndpoints") - self.groupEndpoints.setMinimumSize(QSize(496, 200)) - self.groupEndpoints.setMaximumSize(QSize(496, 200)) + self.groupEndpoints.setMinimumSize(QSize(496, 240)) + self.groupEndpoints.setMaximumSize(QSize(496, 240)) self.formLayout = QFormLayout(self.groupEndpoints) self.formLayout.setObjectName(u"formLayout") self.formLayout.setLabelAlignment(Qt.AlignRight) @@ -174,10 +174,20 @@ def setupUi(self, AuthPanel): self.formLayout.setLayout(4, QFormLayout.ItemRole.FieldRole, self.horizontalLayout) + self.labelCache = QLabel(self.groupEndpoints) + self.labelCache.setObjectName(u"labelCache") + + self.formLayout.setWidget(5, QFormLayout.ItemRole.LabelRole, self.labelCache) + + self.noCacheCheckBox = QCheckBox(self.groupEndpoints) + self.noCacheCheckBox.setObjectName(u"noCacheCheckBox") + + self.formLayout.setWidget(5, QFormLayout.ItemRole.FieldRole, self.noCacheCheckBox) + self.functionsLabel = QLabel(self.groupEndpoints) self.functionsLabel.setObjectName(u"functionsLabel") - self.formLayout.setWidget(5, QFormLayout.ItemRole.LabelRole, self.functionsLabel) + self.formLayout.setWidget(6, QFormLayout.ItemRole.LabelRole, self.functionsLabel) self.selectFuncs = QPushButton(self.groupEndpoints) self.selectFuncs.setObjectName(u"selectFuncs") @@ -185,7 +195,7 @@ def setupUi(self, AuthPanel): self.selectFuncs.setMaximumSize(QSize(296, 32)) self.selectFuncs.setAutoDefault(False) - self.formLayout.setWidget(5, QFormLayout.ItemRole.FieldRole, self.selectFuncs) + self.formLayout.setWidget(6, QFormLayout.ItemRole.FieldRole, self.selectFuncs) self.rootLayout.addWidget(self.groupEndpoints) @@ -254,6 +264,11 @@ def retranslateUi(self, AuthPanel): self.labelScope.setText(QCoreApplication.translate("AuthPanel", u"Scope", None)) self.radioButton_2.setText(QCoreApplication.translate("AuthPanel", u"Public", None)) self.radioButton.setText(QCoreApplication.translate("AuthPanel", u"Private", None)) + self.labelCache.setText(QCoreApplication.translate("AuthPanel", u"Cache", None)) +#if QT_CONFIG(tooltip) + self.noCacheCheckBox.setToolTip(QCoreApplication.translate("AuthPanel", u"Skip cached results and re-run every processing stage.", None)) +#endif // QT_CONFIG(tooltip) + self.noCacheCheckBox.setText(QCoreApplication.translate("AuthPanel", u"Disable cache", None)) self.functionsLabel.setText(QCoreApplication.translate("AuthPanel", u"Functions", None)) self.selectFuncs.setText(QCoreApplication.translate("AuthPanel", u"Select Functions to Upload", None)) self.cancelButton.setText(QCoreApplication.translate("AuthPanel", u"Cancel", None)) diff --git a/reai_toolkit/app/components/tabs/chat_render.py b/reai_toolkit/app/components/tabs/chat_render.py index 0882655..3cc6f19 100644 --- a/reai_toolkit/app/components/tabs/chat_render.py +++ b/reai_toolkit/app/components/tabs/chat_render.py @@ -6,7 +6,8 @@ from __future__ import annotations -from typing import Optional +import re +from typing import Callable, Optional from reai_toolkit.app.services.chat.schema import ChatState, ToolConfirmation @@ -38,12 +39,60 @@ def parse_jump_href(url: str) -> Optional[int]: return None +ENTITY_KINDS = ("FUNCTION", "ANALYSIS", "COLLECTION") + +# The agent writes platform entities the way the Dashboard's mention serializer +# does: a markdown link whose href is the pseudo-URL `FUNCTION_`. Qt keeps +# such an href verbatim, so the link renders but leads nowhere until it is +# rewritten. Tokens also arrive bare in prose, which no renderer linkifies. +_KINDS = "|".join(ENTITY_KINDS) +ENTITY_HREF_RE = re.compile(rf"\]\((?P{_KINDS})_(?P\d+)\)") +# Not preceded by a word char, `[`, `(` or `/`: those mean the token is already +# a link label, an href we left alone, or part of a URL. +BARE_ENTITY_RE = re.compile(rf"(?{_KINDS})_(?P\d+)\b") +CODE_FENCE_RE = re.compile(r"(```.*?```|`[^`\n]*`)", re.DOTALL) + +EntityResolver = Callable[[str, int], Optional[str]] + + +def rewrite_entity_links(markdown: str, resolve: Optional[EntityResolver]) -> str: + """Point `FUNCTION_` pseudo-URLs at somewhere the panel can actually go. + + `resolve(kind, entity_id)` returns the destination, or None to leave the + token alone rather than render a link that cannot be followed. Code spans + and fenced blocks are left verbatim. + """ + if not markdown or resolve is None: + return markdown + + def _target(match: re.Match) -> Optional[str]: + return resolve(match.group("kind"), int(match.group("id"))) + + def _href(match: re.Match) -> str: + target = _target(match) + return match.group(0) if target is None else f"]({target})" + + def _bare(match: re.Match) -> str: + target = _target(match) + return match.group(0) if target is None else f"[{match.group(0)}]({target})" + + out = [] + for chunk in CODE_FENCE_RE.split(markdown): + if chunk.startswith("`"): + out.append(chunk) + else: + out.append(BARE_ENTITY_RE.sub(_bare, ENTITY_HREF_RE.sub(_href, chunk))) + return "".join(out) + + def _function_links(functions) -> str: parts = [f"[{f.name}]({jump_href(f.ea)})" for f in functions if f.name] return "↪ " + " · ".join(parts) if parts else "" -def render_transcript_markdown(state: ChatState) -> str: +def render_transcript_markdown( + state: ChatState, resolve_entity: Optional[EntityResolver] = None +) -> str: """Build a single markdown document from the chat items.""" parts: list[str] = [] for item in state.items: @@ -83,7 +132,7 @@ def render_transcript_markdown(state: ChatState) -> str: if state.run_status == "error" and state.run_error is not None: parts.append(f"> ⚠ **Error:** {state.run_error.message}") - return "\n\n".join(parts) + return rewrite_entity_links("\n\n".join(parts), resolve_entity) def find_pending_confirmation(state: ChatState) -> Optional[ToolConfirmation]: diff --git a/reai_toolkit/app/components/tabs/chat_tab.py b/reai_toolkit/app/components/tabs/chat_tab.py index 814de6d..2c6a3d2 100644 --- a/reai_toolkit/app/components/tabs/chat_tab.py +++ b/reai_toolkit/app/components/tabs/chat_tab.py @@ -24,7 +24,7 @@ render_transcript_markdown, title_case, ) -from reai_toolkit.app.core.qt_compat import QtCore, QtWidgets, Signal, Slot +from reai_toolkit.app.core.qt_compat import QtCore, QtGui, QtWidgets, Signal, Slot from reai_toolkit.app.services.chat.schema import ChatState @@ -234,6 +234,7 @@ def __init__(self, on_closed: Optional[Callable[[], None]] = None) -> None: self.on_select_conversation: Optional[Callable[[str], None]] = None self.on_request_history: Optional[Callable[[], None]] = None self.on_jump: Optional[Callable[[int], None]] = None + self.resolve_entity: Optional[Callable[[str, int], Optional[str]]] = None self.on_stream_event: Optional[Callable] = None self.on_stream_conversation_created: Optional[Callable[[str], None]] = None self.on_stream_error: Optional[Callable[[str], None]] = None @@ -423,7 +424,7 @@ def _flush_render(self) -> None: sb = self._transcript.verticalScrollBar() prev = sb.value() if sb is not None else 0 at_bottom = sb is None or sb.value() >= sb.maximum() - 4 - md = render_transcript_markdown(state) + md = render_transcript_markdown(state, self.resolve_entity) # Rebuilding the whole document resets the scrollbar to the top, so only # do it when the content actually changed, and afterwards keep the reader # where they were: pinned to the bottom if they were following the stream, @@ -560,6 +561,11 @@ def _on_history_item(self, item) -> None: self.on_select_conversation(uuid) def _on_anchor_clicked(self, url) -> None: - ea = parse_jump_href(url.toString()) - if ea is not None and self.on_jump: - self.on_jump(ea) + href = url.toString() + ea = parse_jump_href(href) + if ea is not None: + if self.on_jump: + self.on_jump(ea) + return + if url.scheme() in ("http", "https"): + QtGui.QDesktopServices.openUrl(url) diff --git a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py index f1d9573..88a7278 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py +++ b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py @@ -1,4 +1,5 @@ from logging import Logger +from typing import Optional import ida_funcs import ida_kernwin @@ -207,7 +208,9 @@ def _on_summary_complete( return self._current_summary = response.data if self._decomp_view is not None: - self._decomp_view.set_predicted_name(response.data.predicted_function_name) + self._decomp_view.set_predicted_name( + self._unapplied_predicted_name(ea, response.data.predicted_function_name) + ) self._rerender() def _on_comments_complete( @@ -317,6 +320,19 @@ def _on_rating_result( if self._decomp_view is not None: self._decomp_view.set_rating(None) + def _unapplied_predicted_name( + self, ea: int, name: Optional[str] + ) -> Optional[str]: + if not name: + return None + if self.ai_decomp_service.current_function_name(ea) == name: + return None + return name + + def _hide_predicted_name(self) -> None: + if self._decomp_view is not None: + self._decomp_view.set_predicted_name(None) + def apply_predicted_name(self, name: str) -> None: ea = self._current_func_vaddr if ea is None or not name: @@ -324,6 +340,7 @@ def apply_predicted_name(self, name: str) -> None: if self.ai_decomp_service.update_function_name(ea, name): self.ai_decomp_service.tag_function_as_renamed(name) + self._hide_predicted_name() self.refresh_disassembly_view() return @@ -332,6 +349,7 @@ def apply_predicted_name(self, name: str) -> None: self.show_info_dialog(msg=f"Could not rename this function to '{name}'.") return self.ai_decomp_service.tag_function_as_renamed(final) + self._hide_predicted_name() self.refresh_disassembly_view() def request_rename(self, display_line: int, word: str) -> None: diff --git a/reai_toolkit/app/coordinators/chat_coordinator.py b/reai_toolkit/app/coordinators/chat_coordinator.py index 8284ca2..991b67e 100644 --- a/reai_toolkit/app/coordinators/chat_coordinator.py +++ b/reai_toolkit/app/coordinators/chat_coordinator.py @@ -19,6 +19,7 @@ from libbs.decompilers.ida.compat import execute_read, execute_ui from reai_toolkit.app.app import App +from reai_toolkit.app.components.tabs.chat_render import jump_href from reai_toolkit.app.components.tabs.chat_tab import ChatPanel, ChatStreamWorker from reai_toolkit.app.coordinators.base_coordinator import BaseCoordinator from reai_toolkit.app.factory import DialogFactory @@ -97,6 +98,7 @@ def _wire_panel(self, panel: ChatPanel) -> None: panel.on_select_conversation = self.load_conversation panel.on_request_history = self.request_history panel.on_jump = self.jump_to + panel.resolve_entity = self.entity_href panel.on_stream_event = self.on_stream_event panel.on_stream_conversation_created = self.on_stream_conversation_created panel.on_stream_error = self.on_stream_error @@ -314,6 +316,32 @@ def _work() -> None: threading.Thread(target=_work, daemon=True).start() + def entity_href(self, kind: str, entity_id: int) -> str | None: + """Where a `FUNCTION_` token in the transcript should lead. + + A function in this database is worth more as a local jump than as a + browser tab, so prefer that and fall back to the portal. + """ + portal: str = (self.app.config_service.portal_url or "").rstrip("/") + + if kind == "FUNCTION": + func_map = self.app.netstore_service.get_function_mapping() + ea = func_map.function_map.get(str(entity_id)) if func_map else None + if ea is not None: + return jump_href(int(ea)) + analysis_id = self.app.netstore_service.get_analysis_id() + if portal and analysis_id is not None: + return f"{portal}/analyses/{analysis_id}?view=functions&fn={entity_id}" + return None + + if not portal: + return None + if kind == "ANALYSIS": + return f"{portal}/analyses/{entity_id}" + if kind == "COLLECTION": + return f"{portal}/collections/{entity_id}" + return None + @execute_ui def jump_to(self, ea: int) -> None: try: diff --git a/reai_toolkit/app/interfaces/base_service.py b/reai_toolkit/app/interfaces/base_service.py index 39d3ce7..ea62ddf 100644 --- a/reai_toolkit/app/interfaces/base_service.py +++ b/reai_toolkit/app/interfaces/base_service.py @@ -75,6 +75,11 @@ def update_function_name(ea: int, new_name: str, check_user_flags: bool = False) return ida_name.set_name(ea, new_name, flags) + @staticmethod + @execute_read + def current_function_name(ea: int) -> str: + return ida_name.get_name(ea) or "" + @staticmethod @execute_read def is_protected_user_name(ea: int) -> bool: diff --git a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py index 1560452..decd440 100644 --- a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py +++ b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py @@ -430,7 +430,6 @@ def _queue_decompilation(self, function_id: int) -> tuple[bool, str | None]: with self.yield_api_client(sdk_config=self.sdk_config) as api_client: FunctionsAIDecompilationApi(api_client).create_ai_decompilation( function_id=function_id, - context_aware=True, ) return True, None except ApiException as e: diff --git a/reai_toolkit/app/services/analysis_status/analysis_status.py b/reai_toolkit/app/services/analysis_status/analysis_status.py index f84df96..0c03c83 100644 --- a/reai_toolkit/app/services/analysis_status/analysis_status.py +++ b/reai_toolkit/app/services/analysis_status/analysis_status.py @@ -1,10 +1,15 @@ -import re import threading import time from typing import Any, Callable from loguru import logger -from revengai import AnalysesCoreApi, BaseResponseStatus, Configuration, Logs, StatusOutput +from revengai import ( + AnalysesCoreApi, + BaseResponseStatus, + Configuration, + GetAnalysisLogsOutputBody, + StatusOutput, +) from reai_toolkit.app.core.netstore_service import SimpleNetStore from reai_toolkit.app.core.shared_schema import GenericApiReturn @@ -47,11 +52,10 @@ def _api_get_status(self, analysis_id: int) -> StatusOutput | None: return analysis_status.data - def _api_get_logs(self, analysis_id: int) -> Logs | None: + def _api_get_logs(self, analysis_id: int) -> GetAnalysisLogsOutputBody | None: with self.yield_api_client(sdk_config=self.sdk_config) as api_client: analyses_client = AnalysesCoreApi(api_client) - response = analyses_client.get_analysis_logs(analysis_id) - return response.data + return analyses_client.v3_get_analysis_logs(analysis_id) def _poll_analysis_status(self, stop_event: threading.Event, analysis_id: int) -> None: """ @@ -74,11 +78,10 @@ def _poll_analysis_status(self, stop_event: threading.Event, analysis_id: int) - fn=lambda: self._api_get_logs(analysis_id) ) - logs: Logs | None + logs: GetAnalysisLogsOutputBody | None if get_logs_response.success and (logs := get_logs_response.data): - for line in logs.logs.splitlines(): - line_without_timestamp = re.sub(r"^[\d\-]+ [\d:]+ - ", "", line) - logger.debug(f"RevEng.AI Remote Analysis - {line_without_timestamp}") + for entry in (logs.entries or []): + logger.debug(f"RevEng.AI Remote Analysis - {entry.text}") if get_status_response.data and (status := get_status_response.data.analysis_status): logger.info(f"RevEng.AI: Status - {status}") diff --git a/reai_toolkit/app/services/analysis_sync/analysis_sync.py b/reai_toolkit/app/services/analysis_sync/analysis_sync.py index 9afa121..ef50857 100644 --- a/reai_toolkit/app/services/analysis_sync/analysis_sync.py +++ b/reai_toolkit/app/services/analysis_sync/analysis_sync.py @@ -11,8 +11,8 @@ from loguru import logger from revengai import ( AnalysesCoreApi, + AnalysisBasicInfoOutputBody, BaseResponseAnalysisFunctionMapping, - BaseResponseBasic, Configuration, ) @@ -91,16 +91,18 @@ def _fetch_model_id(self, analysis_id: int) -> int: with self.yield_api_client(sdk_config=self.sdk_config) as api_client: analyses_client = AnalysesCoreApi(api_client) - analysis_details: BaseResponseBasic = analyses_client.get_analysis_basic_info(analysis_id) - model_id = analysis_details.data.model_id + analysis_details: AnalysisBasicInfoOutputBody = ( + analyses_client.get_analysis_basic_info_0(analysis_id) + ) + model_id = analysis_details.model_id self.netstore_service.put_model_id(model_id) - model_name = analysis_details.data.model_name + model_name = analysis_details.model_name self.netstore_service.put_model_name(model_name) local_base_address: int = self._get_current_base_address() - if analysis_details.data and analysis_details.data.base_address is not None: - remote_base_address: int = analysis_details.data.base_address + if analysis_details.base_address is not None: + remote_base_address: int = analysis_details.base_address if local_base_address != remote_base_address: base_address_delta: int = remote_base_address - local_base_address diff --git a/reai_toolkit/app/services/data_types/analysis_catalogue.py b/reai_toolkit/app/services/data_types/analysis_catalogue.py index 137a2ee..fd1fe02 100644 --- a/reai_toolkit/app/services/data_types/analysis_catalogue.py +++ b/reai_toolkit/app/services/data_types/analysis_catalogue.py @@ -7,8 +7,21 @@ from revengai import ( Configuration, CreateAnalysisDataTypesInputBody, + CreateDataTypeEntry, + CreateEnumDataType, + CreateStructDataType, + CreateTypedefDataType, + DataTypeEnumValueEntry, + DataTypeMemberEntry, DataTypesApi, + EnumDefinition, + StructDefinition, + TypedefDefinition, UpdateAnalysisDataTypesInputBody, + UpdateDataTypeEntry, + UpdateEnumDataType, + UpdateStructDataType, + UpdateTypedefDataType, ) from reai_toolkit.app.core.netstore_service import SimpleNetStore @@ -24,6 +37,18 @@ _DEFINED_KINDS = frozenset({KIND_STRUCT, KIND_ENUM, KIND_TYPEDEF}) +_CREATE_MODELS = { + KIND_STRUCT: CreateStructDataType, + KIND_ENUM: CreateEnumDataType, + KIND_TYPEDEF: CreateTypedefDataType, +} + +_UPDATE_MODELS = { + KIND_STRUCT: UpdateStructDataType, + KIND_ENUM: UpdateEnumDataType, + KIND_TYPEDEF: UpdateTypedefDataType, +} + TypeKey = tuple[str, str, str] @@ -31,6 +56,14 @@ def type_key(namespace: Optional[str], name: str, kind: str) -> TypeKey: return (namespace or "", name, (kind or "").upper()) +def empty_definition(kind: str) -> Any: + if kind == KIND_STRUCT: + return StructDefinition(members=[]) + if kind == KIND_ENUM: + return EnumDefinition(values=[]) + return TypedefDefinition(target_data_type_id=None) + + def kind_of(artifact: Any) -> Optional[str]: if isinstance(artifact, Struct): return KIND_STRUCT @@ -153,15 +186,22 @@ def ensure(self, analysis_id: int, artifacts: dict[str, Any]) -> dict[str, int]: def _create_placeholders( self, analysis_id: int, missing: dict[str, Any], catalogue: Catalogue ) -> dict[str, int]: - bodies = [ - { - "kind": kind_of(artifact), - "name": name, - "namespace": "", - "size": getattr(artifact, "size", None) or 0, - } - for name, artifact in missing.items() - ] + bodies = [] + for name, artifact in missing.items(): + kind = kind_of(artifact) + if kind not in _CREATE_MODELS: + continue + bodies.append( + CreateDataTypeEntry( + _CREATE_MODELS[kind]( + kind=kind, + name=name, + namespace="", + size=getattr(artifact, "size", None) or 0, + definition=empty_definition(kind), + ) + ) + ) created: dict[str, int] = {} for chunk in _chunks(bodies, WRITE_BATCH_SIZE): @@ -169,7 +209,7 @@ def _create_placeholders( client = DataTypesApi(api_client) out = client.v3_create_analysis_data_types( analysis_id=analysis_id, - create_analysis_data_types_input_body=CreateAnalysisDataTypesInputBody.model_construct( + create_analysis_data_types_input_body=CreateAnalysisDataTypesInputBody( data_types=chunk ), ).to_dict() @@ -193,18 +233,23 @@ def _write_definitions( ) -> None: bodies = [] for name, artifact in artifacts.items(): + kind = kind_of(artifact) + if kind not in _UPDATE_MODELS: + continue definition = self._definition(artifact, analysis_id, ids, catalogue) if definition is None: continue bodies.append( - { - "data_type_id": ids[name], - "kind": kind_of(artifact), - "name": name, - "namespace": "", - "size": getattr(artifact, "size", None) or 0, - "definition": definition, - } + UpdateDataTypeEntry( + _UPDATE_MODELS[kind]( + data_type_id=ids[name], + kind=kind, + name=name, + namespace="", + size=getattr(artifact, "size", None) or 0, + definition=definition, + ) + ) ) for chunk in _chunks(bodies, WRITE_BATCH_SIZE): @@ -212,7 +257,7 @@ def _write_definitions( client = DataTypesApi(api_client) out = client.v3_update_analysis_data_types( analysis_id=analysis_id, - update_analysis_data_types_input_body=UpdateAnalysisDataTypesInputBody.model_construct( + update_analysis_data_types_input_body=UpdateAnalysisDataTypesInputBody( data_types=chunk ), ).to_dict() @@ -224,31 +269,33 @@ def _definition( analysis_id: int, ids: dict[str, int], catalogue: Catalogue, - ) -> Optional[dict]: + ) -> Optional[Any]: if isinstance(artifact, Struct): members = [] for member in (artifact.members or {}).values(): members.append( - { - "name": member.name or "", - "offset": member.offset, - "size": member.size or 0, - "is_bitfield": False, - "data_type_id": self._reference(member.type, ids, catalogue), - } + DataTypeMemberEntry( + name=member.name or "", + offset=member.offset, + size=member.size or 0, + is_bitfield=False, + data_type_id=self._reference(member.type, ids, catalogue), + ) ) - return {"members": members} + return StructDefinition(members=members) if isinstance(artifact, Enum): - return { - "values": [ - {"name": str(key), "value": str(value)} + return EnumDefinition( + values=[ + DataTypeEnumValueEntry(name=str(key), value=str(value)) for key, value in (artifact.members or {}).items() ] - } + ) if isinstance(artifact, Typedef): - return {"target_data_type_id": self._reference(artifact.type, ids, catalogue)} + return TypedefDefinition( + target_data_type_id=self._reference(artifact.type, ids, catalogue) + ) return None diff --git a/reai_toolkit/app/services/rename/rename_service.py b/reai_toolkit/app/services/rename/rename_service.py index 4bed6b1..2da57ee 100644 --- a/reai_toolkit/app/services/rename/rename_service.py +++ b/reai_toolkit/app/services/rename/rename_service.py @@ -1,19 +1,21 @@ import queue import threading import time +from dataclasses import dataclass from typing import List, Optional from libbs.decompilers.ida.compat import execute_write from loguru import logger from revengai import ( - BaseResponse, + ApiException, + BatchRenameInputBody, + BatchRenameItem, + BatchRenameOutputBody, CanonicalizeNamesInputBody, Configuration, FunctionMapping, - FunctionRenameMap, FunctionsCoreApi, - FunctionsListRename, FunctionsRenamingHistoryApi, ) @@ -25,6 +27,13 @@ from reai_toolkit.app.services.rename.schema import RenameInput +@dataclass +class RemoteRenameOutcome: + status: bool + renamed_count: int = 0 + server_error: bool = False + + class RenameService(IThreadService): _rename_q: queue.Queue[List[RenameInput]] = queue.Queue() _rename_last_ts: dict[int, float] = {} @@ -79,8 +88,14 @@ def _rename_worker(self, stop_event: Optional[threading.Event] = None) -> None: total_errors = None - while attempt < self._rename_max_retries and not stop_event.is_set(): - total_errors = self._rename_function(function_list=function_list) + while attempt < self._rename_max_retries and not ( + stop_event and stop_event.is_set() + ): + try: + total_errors = self._rename_function(function_list=function_list) + except Exception as e: + logger.error(f"RevEng.AI: failed to rename functions: {e}") + total_errors = len(function_list) if total_errors == 0: break attempt += 1 @@ -129,24 +144,24 @@ def _rename_function(self, function_list: List[RenameInput]) -> int: return total_errors # Rename remote functions - response: BaseResponse = self._rename_remote_function( - matched_func_list - ) + response = self._rename_remote_function(matched_func_list) - if not response.status: + if not getattr(response, "status", False): total_errors += len(matched_func_list) return total_errors @execute_write - def _rename_remote_function(self, function_list: list[RenameInput]) -> BaseResponse: - function_rename_list: list[FunctionRenameMap] = [] + def _rename_remote_function( + self, function_list: list[RenameInput] + ) -> RemoteRenameOutcome: + function_rename_list: list[BatchRenameItem] = [] for func in function_list: if func.function_id is None: continue function_rename_list.append( - FunctionRenameMap( + BatchRenameItem( function_id=func.function_id, new_mangled_name=func.new_name, new_name=demangle(func.new_name), @@ -154,25 +169,52 @@ def _rename_remote_function(self, function_list: list[RenameInput]) -> BaseRespo ) self.tag_function_as_renamed(func.new_name) + if not function_rename_list: + return RemoteRenameOutcome(status=True) + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: functions_api = FunctionsRenamingHistoryApi(api_client=api_client) - return functions_api.batch_rename_function( - functions_list_rename=FunctionsListRename( - functions=function_rename_list + try: + response: BatchRenameOutputBody = functions_api.batch_rename_functions( + batch_rename_input_body=BatchRenameInputBody( + functions=function_rename_list + ) ) + except ApiException as e: + logger.error( + f"RevEng.AI: failed to rename {len(function_rename_list)} function(s) " + f"remotely: HTTP {e.status} {e.reason}" + ) + return RemoteRenameOutcome( + status=False, server_error=bool(e.status and e.status >= 500) + ) + except Exception as e: + logger.error(f"RevEng.AI: failed to rename functions remotely: {e}") + return RemoteRenameOutcome(status=False) + + return RemoteRenameOutcome( + status=True, renamed_count=getattr(response, "renamed_count", 0) ) def push_remote_names(self, renames: list[RenameInput]) -> int: pushed: int = 0 for start in range(0, len(renames), self._rename_batch_size): chunk: list[RenameInput] = renames[start:start + self._rename_batch_size] - if getattr(self._rename_remote_function(chunk), "status", False): - pushed += len(chunk) + outcome = self._rename_remote_function(chunk) + if getattr(outcome, "status", False): + pushed += getattr(outcome, "renamed_count", 0) continue + if getattr(outcome, "server_error", False): + logger.error( + f"RevEng.AI: abandoning the push of {len(renames) - pushed} name(s); " + "the platform failed to process the request" + ) + break for rename in chunk: - if getattr(self._rename_remote_function([rename]), "status", False): - pushed += 1 + single = self._rename_remote_function([rename]) + if getattr(single, "status", False): + pushed += getattr(single, "renamed_count", 0) return pushed def canonicalize_names(self, names: list[str]) -> dict[str, str]: diff --git a/reai_toolkit/app/services/upload/upload_service.py b/reai_toolkit/app/services/upload/upload_service.py index ad0b64c..1c06f35 100644 --- a/reai_toolkit/app/services/upload/upload_service.py +++ b/reai_toolkit/app/services/upload/upload_service.py @@ -7,6 +7,7 @@ from revengai import AnalysesCoreApi, BaseResponseConfigResponse, Configuration, Symbols from revengai.api.config_api import ConfigApi +from revengai.models.analysis_config import AnalysisConfig from revengai.models.analysis_create_request import AnalysisCreateRequest from revengai.models.analysis_create_response import AnalysisCreateResponse from revengai.models.analysis_scope import AnalysisScope @@ -36,6 +37,7 @@ def start_analysis( debug_file_path: str | None = None, tags: Optional[list[str]] = None, public: bool = True, + no_cache: bool = False, thread_callback: Optional[Callable[[GenericApiReturn], None]] = None ) -> None: """ @@ -51,6 +53,7 @@ def start_analysis( debug_file_path, tags, public, + no_cache, ), ) @@ -67,6 +70,7 @@ def analyse_file( debug_file_path: str | None = None, tags: Optional[list[str]] = None, public: bool = True, + no_cache: bool = False, ) -> None: """ Uploads the binary (and optional debug file) and starts an analysis. @@ -119,7 +123,8 @@ def analyse_file( debug_sha256=debug_sha256, tags=tags or [], public=public, - symbols=symbols + symbols=symbols, + no_cache=no_cache, ) self.call_callback(generic_return=final_response) @@ -197,6 +202,7 @@ def analyse( debug_sha256: Optional[str] = None, tags: Optional[list[str]] = None, public: bool = True, + no_cache: bool = False, ) -> GenericApiReturn[AnalysisCreateResponse]: if symbols.function_boundaries is None: return GenericApiReturn( @@ -216,6 +222,7 @@ def analyse( tags=[Tag(name=tag) for tag in (tags or [])], analysis_scope=AnalysisScope.PUBLIC if public else AnalysisScope.PRIVATE, symbols=symbols, + analysis_config=AnalysisConfig(no_cache=no_cache), ) response = self.api_request_returning( diff --git a/tests/unit/analysis_status/test_sdk_schemas.py b/tests/unit/analysis_status/test_sdk_schemas.py index 7b9f64d..c9ef268 100644 --- a/tests/unit/analysis_status/test_sdk_schemas.py +++ b/tests/unit/analysis_status/test_sdk_schemas.py @@ -1,12 +1,21 @@ -from revengai import BaseResponseStatus, Logs, StatusOutput +from revengai import ( + AnalysisLogEntry, + BaseResponseStatus, + GetAnalysisLogsOutputBody, + StatusOutput, +) def test_status_output_fields(): assert {"analysis_id", "analysis_status"} <= set(StatusOutput.model_fields) -def test_logs_field(): - assert "logs" in Logs.model_fields +def test_logs_output_carries_entries(): + assert "entries" in GetAnalysisLogsOutputBody.model_fields + + +def test_log_entry_fields(): + assert {"text", "level", "timestamp"} <= set(AnalysisLogEntry.model_fields) def test_base_response_status_envelope(): diff --git a/tests/unit/analysis_status/test_service.py b/tests/unit/analysis_status/test_service.py index 1f15c26..8cbdb23 100644 --- a/tests/unit/analysis_status/test_service.py +++ b/tests/unit/analysis_status/test_service.py @@ -3,7 +3,7 @@ import pytest from revengai import ApiException -from revengai.models.logs import Logs +from revengai.models.get_analysis_logs_output_body import GetAnalysisLogsOutputBody from revengai.models.status_output import StatusOutput from reai_toolkit.app.services.analysis_status import analysis_status as svc_mod @@ -35,7 +35,9 @@ def sdk(mocker): api_class = mocker.patch.object(svc_mod, "AnalysesCoreApi") api_inst = MagicMock() api_class.return_value = api_inst - api_inst.get_analysis_logs.return_value = MagicMock(data=Logs.model_construct(logs="")) + api_inst.v3_get_analysis_logs.return_value = GetAnalysisLogsOutputBody.model_construct( + entries=[] + ) return api_inst @@ -81,7 +83,7 @@ def test_status_request_failure_surfaces_error(service, sdk): result = cb.call_args[0][0] assert result.success is False - sdk.get_analysis_logs.assert_not_called() + sdk.v3_get_analysis_logs.assert_not_called() def test_in_progress_status_keeps_polling_until_complete(service, sdk): diff --git a/tests/unit/analysis_sync/test_sdk_schemas.py b/tests/unit/analysis_sync/test_sdk_schemas.py index a290b96..8e172b8 100644 --- a/tests/unit/analysis_sync/test_sdk_schemas.py +++ b/tests/unit/analysis_sync/test_sdk_schemas.py @@ -1,13 +1,15 @@ from revengai import ( + AnalysisBasicInfoOutputBody, BaseResponseAnalysisFunctionMapping, - BaseResponseBasic, FunctionMapping, ) from revengai.models.analysis_function_mapping import AnalysisFunctionMapping -def test_basic_response_envelope(): - assert {"status", "data", "errors"} <= set(BaseResponseBasic.model_fields) +def test_basic_info_carries_model_and_base_address(): + assert {"model_id", "model_name", "base_address"} <= set( + AnalysisBasicInfoOutputBody.model_fields + ) def test_function_mapping_envelope_carries_function_maps(): diff --git a/tests/unit/analysis_sync/test_service.py b/tests/unit/analysis_sync/test_service.py index 6af98ad..35c2f0e 100644 --- a/tests/unit/analysis_sync/test_service.py +++ b/tests/unit/analysis_sync/test_service.py @@ -1,8 +1,12 @@ from unittest.mock import MagicMock import pytest +from revengai.models.analysis_basic_info_output_body import ( + AnalysisBasicInfoOutputBody, +) from revengai.models.function_mapping import FunctionMapping +from reai_toolkit.app.services.analysis_sync import analysis_sync as svc_mod from reai_toolkit.app.services.analysis_sync.analysis_sync import AnalysisSyncService from reai_toolkit.app.services.analysis_sync.schema import MatchedFunctionSummary from reai_toolkit.app.services.data_types.data_types_service import DataTypesImportResult @@ -208,3 +212,57 @@ def test_get_function_matches_aborts_when_model_id_fails(service, mocker): fetch_map.assert_not_called() cb.assert_not_called() + + +def _basic_info(model_id=7, base_address=0x140000000): + return AnalysisBasicInfoOutputBody.model_construct( + model_id=model_id, model_name="binnet", base_address=base_address + ) + + +@pytest.fixture +def analyses_api(mocker): + mocker.patch.object(AnalysisSyncService, "yield_api_client") + api_class = mocker.patch.object(svc_mod, "AnalysesCoreApi") + api = MagicMock() + api_class.return_value = api + return api + + +def test_fetch_model_id_reads_v3_basic_info(service, analyses_api, netstore, mocker): + mocker.patch.object(service, "_get_current_base_address", return_value=0x140000000) + rebase = mocker.patch.object(service, "_rebase_program") + analyses_api.get_analysis_basic_info_0.return_value = _basic_info() + + assert service._fetch_model_id(1234) == 7 + + analyses_api.get_analysis_basic_info_0.assert_called_once_with(1234) + netstore.put_model_id.assert_called_once_with(7) + netstore.put_model_name.assert_called_once_with("binnet") + rebase.assert_not_called() + + +def test_fetch_model_id_rebases_on_base_address_mismatch( + service, analyses_api, mocker +): + mocker.patch.object(service, "_get_current_base_address", return_value=0x400000) + rebase = mocker.patch.object(service, "_rebase_program") + analyses_api.get_analysis_basic_info_0.return_value = _basic_info( + base_address=0x140000000 + ) + + service._fetch_model_id(1234) + + rebase.assert_called_once_with(0x140000000 - 0x400000) + + +def test_fetch_model_id_skips_rebase_when_base_address_unknown( + service, analyses_api, mocker +): + mocker.patch.object(service, "_get_current_base_address", return_value=0x400000) + rebase = mocker.patch.object(service, "_rebase_program") + analyses_api.get_analysis_basic_info_0.return_value = _basic_info(base_address=None) + + service._fetch_model_id(1234) + + rebase.assert_not_called() diff --git a/tests/unit/chat/test_render.py b/tests/unit/chat/test_render.py index d82dc83..e1160df 100644 --- a/tests/unit/chat/test_render.py +++ b/tests/unit/chat/test_render.py @@ -12,6 +12,7 @@ jump_href, parse_jump_href, render_transcript_markdown, + rewrite_entity_links, title_case, ) from reai_toolkit.app.services.chat.reducer import build_initial_state @@ -123,3 +124,65 @@ def test_find_pending_confirmation(): pending = find_pending_confirmation(state) assert pending is not None and pending.id == "c2" assert find_pending_confirmation(ChatState()) is None + + +def _resolver(kind, entity_id): + if kind == "FUNCTION" and entity_id == 1: + return jump_href(0x140001000) + if kind == "FUNCTION": + return f"https://portal.reveng.ai/analyses/9?view=functions&fn={entity_id}" + if kind == "ANALYSIS": + return f"https://portal.reveng.ai/analyses/{entity_id}" + return None + + +def test_function_href_in_this_database_becomes_a_local_jump(): + assert rewrite_entity_links("[foo](FUNCTION_1)", _resolver) == ( + f"[foo]({jump_href(0x140001000)})" + ) + + +def test_function_href_elsewhere_falls_back_to_the_portal(): + assert rewrite_entity_links("[foo](FUNCTION_42)", _resolver) == ( + "[foo](https://portal.reveng.ai/analyses/9?view=functions&fn=42)" + ) + + +def test_bare_token_in_prose_becomes_a_link(): + assert rewrite_entity_links("renamed FUNCTION_42 today", _resolver) == ( + "renamed [FUNCTION_42]" + "(https://portal.reveng.ai/analyses/9?view=functions&fn=42) today" + ) + + +def test_analysis_token_points_at_the_portal(): + assert rewrite_entity_links("[run](ANALYSIS_7)", _resolver) == ( + "[run](https://portal.reveng.ai/analyses/7)" + ) + + +def test_unresolvable_token_is_left_alone_rather_than_linked_nowhere(): + text = "see [c](COLLECTION_5) and COLLECTION_5" + assert rewrite_entity_links(text, _resolver) == text + + +def test_tokens_inside_code_are_left_verbatim(): + fenced = "```\nundefined8 FUNCTION_42(void)\n```" + assert rewrite_entity_links(fenced, _resolver) == fenced + assert rewrite_entity_links("`FUNCTION_42`", _resolver) == "`FUNCTION_42`" + + +def test_a_token_used_as_link_text_is_not_double_wrapped(): + text = f"[FUNCTION_42]({jump_href(16)})" + assert rewrite_entity_links(text, _resolver) == text + + +def test_render_without_a_resolver_leaves_the_transcript_untouched(): + state = ChatState(items=[AssistantMessage(id="m1", content="[foo](FUNCTION_1)", is_streaming=False)]) + assert "FUNCTION_1" in render_transcript_markdown(state) + + +def test_render_applies_the_resolver_to_assistant_text(): + state = ChatState(items=[AssistantMessage(id="m1", content="[foo](FUNCTION_1)", is_streaming=False)]) + out = render_transcript_markdown(state, _resolver) + assert jump_href(0x140001000) in out diff --git a/tests/unit/core/test_sdk_call_signatures.py b/tests/unit/core/test_sdk_call_signatures.py new file mode 100644 index 0000000..92d408a --- /dev/null +++ b/tests/unit/core/test_sdk_call_signatures.py @@ -0,0 +1,205 @@ +""" +Guard against revengai SDK drift. + +The service tests mock the SDK's API classes with plain MagicMocks, which accept +any keyword and return an object with any attribute. A parameter or field the SDK +drops therefore passes the suite and only fails at runtime — as `context_aware` +did when the SDK went from 3.x to 4.x. These checks read the real signatures and +model fields instead. + +Both checks resolve the receiver of a call back to an SDK API class, so a plugin +method that merely shares a name with an SDK one (netstore_service.get_analysis_status, +matching_service.search_binaries) is neither falsely flagged nor able to mask a +real call. +""" + +import ast +import importlib +import inspect +import pkgutil +import typing +from pathlib import Path + +import pydantic +import revengai.api as sdk_api + +PLUGIN_ROOT = Path(__file__).resolve().parents[3] / "reai_toolkit" + +# Attributes every generated model carries that are not declared fields. +MODEL_EXTRA = { + "to_dict", "to_json", "to_str", "from_json", "from_dict", "model_dump", + "model_dump_json", "model_fields", "model_construct", "model_copy", + "additional_properties", "model_fields_set", +} + + +def _api_classes() -> dict[str, type]: + classes: dict[str, type] = {} + for module in pkgutil.iter_modules(sdk_api.__path__): + mod = importlib.import_module(f"revengai.api.{module.name}") + for name, cls in vars(mod).items(): + if inspect.isclass(cls) and name.endswith("Api"): + classes[name] = cls + return classes + + +API_CLASSES = _api_classes() + + +def _signature(api_name: str, method: str): + cls = API_CLASSES.get(api_name) + fn = getattr(cls, method, None) if cls else None + if fn is None or not callable(fn): + return None + try: + return inspect.signature(fn) + except (TypeError, ValueError): + return None + + +def _model_in(annotation): + """Pull the model class out of Optional[X], List[X], Union[...].""" + seen, stack = set(), [annotation] + while stack: + candidate = stack.pop() + if inspect.isclass(candidate) and issubclass(candidate, pydantic.BaseModel): + return candidate + if id(candidate) in seen: + continue + seen.add(id(candidate)) + stack.extend(typing.get_args(candidate)) + return None + + +def _chain(node: ast.Attribute) -> tuple[ast.AST, list[str]]: + parts = [] + while isinstance(node, ast.Attribute): + parts.append(node.attr) + node = node.value + return node, list(reversed(parts)) + + +def _constructed_api(node) -> str | None: + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): + if node.func.id in API_CLASSES: + return node.func.id + return None + + +def _plugin_functions(): + for path in sorted(PLUGIN_ROOT.rglob("*.py")): + if "vendor" in path.parts: + continue + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + yield path, node + + +def _sdk_calls(fn_node): + """Every call in this function that lands on an SDK API class.""" + clients, rebound = {}, set() + for node in ast.walk(fn_node): + if not isinstance(node, (ast.Assign, ast.AnnAssign)): + continue + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + for target in targets: + if not isinstance(target, ast.Name): + continue + api = _constructed_api(node.value) + if api: + clients[target.id] = api + else: + rebound.add(target.id) + + def receiver(node) -> str | None: + if isinstance(node, ast.Name) and node.id in clients and node.id not in rebound: + return clients[node.id] + return _constructed_api(node) + + calls = {} + for node in ast.walk(fn_node): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)): + continue + api = receiver(node.func.value) + if not api: + continue + signature = _signature(api, node.func.attr) + if signature is not None: + calls[node] = (api, node.func.attr, signature) + return calls + + +def _relative(path: Path) -> str: + return str(path.relative_to(PLUGIN_ROOT.parent)) + + +def test_every_keyword_the_plugin_passes_is_an_sdk_parameter(): + offenders = [] + for path, fn_node in _plugin_functions(): + for call, (api, method, signature) in _sdk_calls(fn_node).items(): + for keyword in call.keywords: + if keyword.arg and keyword.arg not in signature.parameters: + offenders.append( + f"{_relative(path)}:{call.lineno} " + f"{api}.{method}({keyword.arg}=...) is not a parameter" + ) + assert sorted(set(offenders)) == [] + + +def test_every_attribute_read_off_an_sdk_response_is_a_model_field(): + offenders = [] + for path, fn_node in _plugin_functions(): + calls = _sdk_calls(fn_node) + + typed, rebound = {}, set() + for node in ast.walk(fn_node): + target = value = None + if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): + target, value = node.target.id, node.value + elif isinstance(node, ast.Assign) and len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + target, value = node.targets[0].id, node.value + elif isinstance(node, ast.NamedExpr) and isinstance(node.target, ast.Name): + target, value = node.target.id, node.value + if target is None: + continue + call = calls.get(value) + model = _model_in(call[2].return_annotation) if call else None + if model is None: + rebound.add(target) + continue + if target in typed and typed[target][1] is not model: + rebound.add(target) + typed[target] = (f"{call[0]}.{call[1]}", model) + + for node in ast.walk(fn_node): + if not isinstance(node, ast.Attribute): + continue + base, parts = _chain(node) + origin = label = None + if isinstance(base, ast.Name) and base.id not in rebound: + origin, label = typed.get(base.id), base.id + elif isinstance(base, ast.Call) and base in calls: + api, method, signature = calls[base] + model = _model_in(signature.return_annotation) + if model: + origin, label = (f"{api}.{method}", model), f"{method}(...)" + if not origin: + continue + + source, model = origin + walked = label + for part in parts: + if model is None or part in MODEL_EXTRA or part.startswith("_"): + break + if part not in model.model_fields: + offenders.append( + f"{_relative(path)}:{node.lineno} {walked}.{part} — " + f"{source}() returns {model.__name__}, which has no '{part}' " + f"(fields: {', '.join(sorted(model.model_fields))})" + ) + break + walked = f"{walked}.{part}" + model = _model_in(model.model_fields[part].annotation) + + assert sorted(set(offenders)) == [] diff --git a/tests/unit/data_types/test_analysis_catalogue.py b/tests/unit/data_types/test_analysis_catalogue.py index 8159d36..97941be 100644 --- a/tests/unit/data_types/test_analysis_catalogue.py +++ b/tests/unit/data_types/test_analysis_catalogue.py @@ -51,6 +51,10 @@ def _written(items): return resp +def _sent(body): + return body.to_dict()["data_types"] + + def test_catalogue_indexes_by_key_and_name(): catalogue = Catalogue.of([_entry(1, "Foo"), _entry(2, "Bar", kind="ENUM")]) @@ -128,7 +132,7 @@ def test_ensure_creates_only_the_gaps(service, sdk): body = sdk.v3_create_analysis_data_types.call_args.kwargs[ "create_analysis_data_types_input_body" ] - assert [d["name"] for d in body.data_types] == ["New"] + assert [d["name"] for d in _sent(body)] == ["New"] def test_create_bodies_carry_no_data_type_id(service, sdk): @@ -141,7 +145,7 @@ def test_create_bodies_carry_no_data_type_id(service, sdk): body = sdk.v3_create_analysis_data_types.call_args.kwargs[ "create_analysis_data_types_input_body" ] - assert all("data_type_id" not in d for d in body.data_types) + assert all("data_type_id" not in d for d in _sent(body)) def test_definitions_are_written_after_ids_exist(service, sdk): @@ -163,10 +167,10 @@ def test_definitions_are_written_after_ids_exist(service, sdk): body = sdk.v3_update_analysis_data_types.call_args.kwargs[ "update_analysis_data_types_input_body" ] - written = {d["name"]: d for d in body.data_types} + written = {d["name"]: d for d in _sent(body)} assert written["Node"]["definition"]["members"][0]["data_type_id"] == 1 assert written["NodeRef"]["definition"]["target_data_type_id"] == 2 - assert all("data_type_id" in d for d in body.data_types) + assert all("data_type_id" in d for d in _sent(body)) def test_self_referential_struct_resolves_to_its_own_new_id(service, sdk): @@ -185,7 +189,7 @@ def test_self_referential_struct_resolves_to_its_own_new_id(service, sdk): body = sdk.v3_update_analysis_data_types.call_args.kwargs[ "update_analysis_data_types_input_body" ] - assert body.data_types[0]["definition"]["members"][0]["data_type_id"] == 5 + assert _sent(body)[0]["definition"]["members"][0]["data_type_id"] == 5 def test_unresolvable_member_type_is_left_unresolved(service, sdk): @@ -204,7 +208,7 @@ def test_unresolvable_member_type_is_left_unresolved(service, sdk): body = sdk.v3_update_analysis_data_types.call_args.kwargs[ "update_analysis_data_types_input_body" ] - assert body.data_types[0]["definition"]["members"][0]["data_type_id"] is None + assert "data_type_id" not in _sent(body)[0]["definition"]["members"][0] def test_enum_values_are_sent_as_strings(service, sdk): @@ -219,7 +223,7 @@ def test_enum_values_are_sent_as_strings(service, sdk): body = sdk.v3_update_analysis_data_types.call_args.kwargs[ "update_analysis_data_types_input_body" ] - values = body.data_types[0]["definition"]["values"] + values = _sent(body)[0]["definition"]["values"] assert {v["name"]: v["value"] for v in values} == { "RED": "0", "BIG": "18446744073709551615", @@ -233,7 +237,7 @@ def test_writes_are_chunked(service, sdk): [ _entry(i, d["name"]) for i, d in enumerate( - kw["create_analysis_data_types_input_body"].data_types, start=1 + _sent(kw["create_analysis_data_types_input_body"]), start=1 ) ] ) @@ -245,7 +249,7 @@ def test_writes_are_chunked(service, sdk): service.ensure(ANALYSIS, artifacts) sizes = [ - len(c.kwargs["create_analysis_data_types_input_body"].data_types) + len(_sent(c.kwargs["create_analysis_data_types_input_body"])) for c in sdk.v3_create_analysis_data_types.call_args_list ] assert sizes == [WRITE_BATCH_SIZE, 10] diff --git a/tests/unit/rename/test_sdk_schemas.py b/tests/unit/rename/test_sdk_schemas.py index 274e173..93c7274 100644 --- a/tests/unit/rename/test_sdk_schemas.py +++ b/tests/unit/rename/test_sdk_schemas.py @@ -1,19 +1,19 @@ from revengai import ( - BaseResponse, - FunctionRenameMap, - FunctionsListRename, + BatchRenameInputBody, + BatchRenameItem, + BatchRenameOutputBody, ) -def test_function_rename_map_fields(): +def test_batch_rename_item_fields(): assert {"function_id", "new_name", "new_mangled_name"} <= set( - FunctionRenameMap.model_fields + BatchRenameItem.model_fields ) -def test_functions_list_rename_wraps_functions(): - assert "functions" in FunctionsListRename.model_fields +def test_batch_rename_input_body_wraps_functions(): + assert "functions" in BatchRenameInputBody.model_fields -def test_base_response_envelope(): - assert {"status", "data"} <= set(BaseResponse.model_fields) +def test_batch_rename_output_reports_renamed_count(): + assert "renamed_count" in BatchRenameOutputBody.model_fields diff --git a/tests/unit/rename/test_service.py b/tests/unit/rename/test_service.py index 529ea83..49c6809 100644 --- a/tests/unit/rename/test_service.py +++ b/tests/unit/rename/test_service.py @@ -3,10 +3,12 @@ from unittest.mock import MagicMock import pytest -from revengai import BaseResponse from revengai.models.function_mapping import FunctionMapping -from reai_toolkit.app.services.rename.rename_service import RenameService +from reai_toolkit.app.services.rename.rename_service import ( + RemoteRenameOutcome, + RenameService, +) from reai_toolkit.app.services.rename.schema import RenameInput @@ -29,7 +31,7 @@ def ida_calls(mocker): remote = mocker.patch.object( RenameService, "_rename_remote_function", - return_value=BaseResponse.model_construct(status=True), + return_value=RemoteRenameOutcome(status=True, renamed_count=1), ) return update, remote @@ -79,7 +81,7 @@ def test_rename_function_counts_local_failure_and_skips_remote(service, ida_call def test_rename_function_counts_remote_failure(service, ida_calls): _, remote = ida_calls - remote.return_value = BaseResponse.model_construct(status=False) + remote.return_value = RemoteRenameOutcome(status=False) errors = service._rename_function( [RenameInput(ea=0x10, new_name="foo", function_id=1)] @@ -88,6 +90,28 @@ def test_rename_function_counts_remote_failure(service, ida_calls): assert errors == 1 +def test_rename_function_counts_absent_remote_response(service, ida_calls): + _, remote = ida_calls + remote.return_value = None + + errors = service._rename_function( + [RenameInput(ea=0x10, new_name="foo", function_id=1)] + ) + + assert errors == 1 + + +def test_push_remote_names_stops_on_server_error(service, ida_calls): + _, remote = ida_calls + renames = [RenameInput(ea=i, new_name=f"f{i}", function_id=i) for i in range(3)] + remote.side_effect = [RemoteRenameOutcome(status=False, server_error=True)] + + pushed = service.push_remote_names(renames) + + assert pushed == 0 + assert remote.call_count == 1 + + def test_rename_function_resolves_missing_function_id_from_mapping( service, ida_calls, netstore ): @@ -126,8 +150,8 @@ def test_push_remote_names_delegates_to_remote_rename(service, ida_calls): def test_push_remote_names_falls_back_to_per_item_on_batch_rejection(service, ida_calls): _, remote = ida_calls renames = [RenameInput(ea=i, new_name=f"f{i}", function_id=i) for i in range(3)] - ok = BaseResponse.model_construct(status=True) - bad = BaseResponse.model_construct(status=False) + ok = RemoteRenameOutcome(status=True, renamed_count=1) + bad = RemoteRenameOutcome(status=False) remote.side_effect = [bad, ok, bad, ok] pushed = service.push_remote_names(renames) diff --git a/tests/unit/upload/test_service.py b/tests/unit/upload/test_service.py index d1d3afe..7522a93 100644 --- a/tests/unit/upload/test_service.py +++ b/tests/unit/upload/test_service.py @@ -89,3 +89,27 @@ def test_analyse_converts_tags_to_tag_models(service, sdk, symbols, tags, expect request = sdk.create_analysis.call_args.kwargs["analysis_create_request"] assert request.tags == [Tag(name=name) for name in expected] + + +@pytest.mark.parametrize("no_cache", [True, False]) +def test_analyse_sets_no_cache_on_analysis_config(service, sdk, symbols, no_cache): + service.analyse( + file_name="binary.elf", + binary_sha256="deadbeef", + symbols=symbols, + no_cache=no_cache, + ) + + request = sdk.create_analysis.call_args.kwargs["analysis_create_request"] + assert request.analysis_config.no_cache is no_cache + + +def test_analyse_defaults_no_cache_false(service, sdk, symbols): + service.analyse( + file_name="binary.elf", + binary_sha256="deadbeef", + symbols=symbols, + ) + + request = sdk.create_analysis.call_args.kwargs["analysis_create_request"] + assert request.analysis_config.no_cache is False