From af04f90049f5f66d35ec739f85f3da9c79471090 Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:19:53 +0100 Subject: [PATCH 1/3] feat(PLU-314): syanc analysis updates --- pyproject.toml | 2 +- reai_toolkit/app/app.py | 11 +- .../coordinators/sync_analysis_coordinator.py | 18 ++- reai_toolkit/app/interfaces/base_service.py | 32 ++++- .../services/analysis_sync/analysis_sync.py | 133 ++++++++++++++---- .../app/services/analysis_sync/schema.py | 4 + .../services/data_types/data_types_service.py | 35 ++++- .../app/services/rename/rename_service.py | 28 ++++ .../variable_sync/variable_sync_service.py | 83 +++++++++++ .../app/transformations/import_data_types.py | 38 +++-- tests/idalib/analysis_sync/test_sync.py | 62 +++++++- tests/unit/analysis_sync/test_service.py | 119 ++++++++++++++-- tests/unit/data_types/test_service.py | 60 +++++--- tests/unit/rename/test_service.py | 51 +++++++ tests/unit/variable_sync/test_service.py | 77 ++++++++++ uv.lock | 8 +- 16 files changed, 665 insertions(+), 96 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5249962..5158749 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ dependencies = [ "requests>=2.32", "loguru>=0.7", "pydantic", - "revengai>=3.110.0", + "revengai>=3.113.0", "libbs>=2.16.5", ] diff --git a/reai_toolkit/app/app.py b/reai_toolkit/app/app.py index a9b8dd4..a2088ed 100644 --- a/reai_toolkit/app/app.py +++ b/reai_toolkit/app/app.py @@ -44,16 +44,19 @@ def __init__(self, ida_version: str = "UNKNOWN", plugin_version: str = "UNKNOWN" self.data_types_service = ImportDataTypesService( netstore_service=self.netstore_service, sdk_config=sdk_config ) - self.analysis_sync_service = AnalysisSyncService(data_types_service=self.data_types_service, + self.rename_service = RenameService( netstore_service=self.netstore_service, sdk_config=sdk_config ) - self.existing_analyses_service = ExistingAnalysesService( + self.variable_sync_service = VariableSyncService( netstore_service=self.netstore_service, sdk_config=sdk_config ) - self.rename_service = RenameService( + self.analysis_sync_service = AnalysisSyncService( + data_types_service=self.data_types_service, + rename_service=self.rename_service, + variable_sync_service=self.variable_sync_service, netstore_service=self.netstore_service, sdk_config=sdk_config ) - self.variable_sync_service = VariableSyncService( + self.existing_analyses_service = ExistingAnalysesService( netstore_service=self.netstore_service, sdk_config=sdk_config ) self.ai_decomp_service = AiDecompService( diff --git a/reai_toolkit/app/coordinators/sync_analysis_coordinator.py b/reai_toolkit/app/coordinators/sync_analysis_coordinator.py index 3d94be3..0c08650 100644 --- a/reai_toolkit/app/coordinators/sync_analysis_coordinator.py +++ b/reai_toolkit/app/coordinators/sync_analysis_coordinator.py @@ -99,10 +99,22 @@ def _on_complete( elif generic_return.data is not None and generic_return.data.data_types_error: self.show_error_dialog(message=generic_return.data.data_types_error) else: - self.show_info_dialog( - msg=f"Analysis data synced successfully. \n\nSynced {generic_return.data.matched_function_count} functions with remote analysis." - + f"\n{generic_return.data.unmatched_function_count} local functions not present in remote analysis." + data = generic_return.data + msg: str = ( + f"Analysis data synced successfully. \n\nSynced {data.matched_function_count} functions with remote analysis." + + f"\n{data.unmatched_function_count} local functions not present in remote analysis." ) + if data.canonicalized_name_count or data.deduped_name_count: + msg += ( + f"\nCanonicalized {data.canonicalized_name_count} and disambiguated " + f"{data.deduped_name_count} name(s)." + ) + if data.pushed_name_count or data.pushed_type_count: + msg += ( + f"\nPushed {data.pushed_name_count} name(s) and {data.pushed_type_count} " + f"type set(s) back to the platform." + ) + self.show_info_dialog(msg=msg) self.refresh_disassembly_view() diff --git a/reai_toolkit/app/interfaces/base_service.py b/reai_toolkit/app/interfaces/base_service.py index 986aa7e..23c9609 100644 --- a/reai_toolkit/app/interfaces/base_service.py +++ b/reai_toolkit/app/interfaces/base_service.py @@ -3,7 +3,8 @@ import ida_name import ida_dirtree -from libbs.decompilers.ida.compat import execute_write +import idaapi +from libbs.decompilers.ida.compat import execute_write, execute_read from loguru import logger from revengai import ApiClient, ApiException, Configuration @@ -74,6 +75,35 @@ 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 is_protected_user_name(ea: int) -> bool: + cur: str = ida_name.get_name(ea) or "" + try: + return bool(cur and ida_name.is_uname(cur)) + except Exception: + return False + + @staticmethod + @execute_write + def apply_deduped_name(ea: int, desired_name: str) -> str | None: + flags: int = ida_name.SN_CHECK | ida_name.SN_AUTO | ida_name.SN_NODUMMY + + candidate: str = desired_name + suffix: int = 1 + while True: + holder: int = ida_name.get_name_ea(idaapi.BADADDR, candidate) + if holder == idaapi.BADADDR or holder == ea: + break + candidate = f"{desired_name}_{suffix}" + suffix += 1 + if suffix > 100000: + return None + + if ida_name.set_name(ea, candidate, flags): + return ida_name.get_name(ea) or candidate + return None + # ========================================= # RevEng API WRAPPER SAFE METHODS # ========================================= diff --git a/reai_toolkit/app/services/analysis_sync/analysis_sync.py b/reai_toolkit/app/services/analysis_sync/analysis_sync.py index 260b59d..571a638 100644 --- a/reai_toolkit/app/services/analysis_sync/analysis_sync.py +++ b/reai_toolkit/app/services/analysis_sync/analysis_sync.py @@ -6,6 +6,7 @@ import idautils import idaapi +import ida_name from loguru import logger from revengai import ( @@ -20,14 +21,26 @@ from reai_toolkit.app.interfaces.thread_service import IThreadService from reai_toolkit.app.services.analysis_sync.schema import MatchedFunctionSummary from reai_toolkit.app.services.data_types.data_types_service import ImportDataTypesService +from reai_toolkit.app.services.rename.rename_service import RenameService +from reai_toolkit.app.services.rename.schema import RenameInput +from reai_toolkit.app.services.variable_sync.variable_sync_service import VariableSyncService from revengai.models.function_mapping import FunctionMapping class AnalysisSyncService(IThreadService): - def __init__(self, data_types_service: ImportDataTypesService, netstore_service: SimpleNetStore, sdk_config: Configuration): + def __init__( + self, + data_types_service: ImportDataTypesService, + rename_service: RenameService, + variable_sync_service: VariableSyncService, + netstore_service: SimpleNetStore, + sdk_config: Configuration, + ): super().__init__(netstore_service=netstore_service, sdk_config=sdk_config) self.data_types_service: ImportDataTypesService = data_types_service + self.rename_service: RenameService = rename_service + self.variable_sync_service: VariableSyncService = variable_sync_service def thread_in_progress(self) -> bool: """ @@ -121,73 +134,131 @@ def _fetch_function_map(self, analysis_id: int) -> FunctionMapping: def _perform_function_sync( self, func_map: FunctionMapping, - ) -> GenericApiReturn[MatchedFunctionSummary]: - # Mapping of local function addresses to mangled names + ) -> tuple[GenericApiReturn[MatchedFunctionSummary], list[RenameInput], list[tuple[int, int | None, str]]]: local_vaddr_to_matched_name: dict[str, str] = func_map.name_map + inverse_map: dict[str, int] = func_map.inverse_function_map logger.info(f"RevEng.AI: Retrieved {len(local_vaddr_to_matched_name)} functions from analysis") - # Compute which IDA functions match the revengai analysis functions matched_function_count: int = 0 unmatched_function_count: int = 0 total_function_count: int = 0 + deduped_name_count: int = 0 + + name_pushbacks: list[RenameInput] = [] + needs_canonical: list[tuple[int, int | None, str]] = [] local_vaddr: int for local_vaddr in idautils.Functions(): local_vaddr_str: str = str(local_vaddr) new_name: str | None = local_vaddr_to_matched_name.get(local_vaddr_str) old_name: str | None = idc.get_func_name(local_vaddr) - if new_name: - if new_name != old_name: - self.update_function_name(local_vaddr, new_name, check_user_flags=True) - self.tag_function_as_renamed(new_name) + if new_name: matched_function_count += 1 + if new_name != old_name: + fid: int | None = inverse_map.get(local_vaddr_str) + if self.update_function_name(local_vaddr, new_name, check_user_flags=True): + self.tag_function_as_renamed(new_name) + elif self.is_protected_user_name(local_vaddr): + pass + else: + holder: int = ida_name.get_name_ea(idaapi.BADADDR, new_name) + if holder != idaapi.BADADDR and holder != local_vaddr: + final: str | None = self.apply_deduped_name(local_vaddr, new_name) + if final: + deduped_name_count += 1 + self.tag_function_as_renamed(final) + if fid is not None: + name_pushbacks.append( + RenameInput(ea=local_vaddr, new_name=final, function_id=fid) + ) + else: + needs_canonical.append((local_vaddr, fid, new_name)) else: unmatched_function_count += 1 - + total_function_count += 1 logger.info(f"RevEng.AI: Matched {matched_function_count} functions") logger.info(f"RevEng.AI: {unmatched_function_count} functions not matched") - return GenericApiReturn( - success=True, - data=MatchedFunctionSummary( - matched_function_count=matched_function_count, - unmatched_function_count=unmatched_function_count, - total_function_count=total_function_count, - ), + summary: MatchedFunctionSummary = MatchedFunctionSummary( + matched_function_count=matched_function_count, + unmatched_function_count=unmatched_function_count, + total_function_count=total_function_count, + deduped_name_count=deduped_name_count, ) + return GenericApiReturn(success=True, data=summary), name_pushbacks, needs_canonical def _match_functions( self, func_map: FunctionMapping - ) -> GenericApiReturn[MatchedFunctionSummary]: - + ) -> tuple[GenericApiReturn[MatchedFunctionSummary], list[RenameInput], list[tuple[int, int | None, str]]]: try: - data: GenericApiReturn[MatchedFunctionSummary] = self._perform_function_sync(func_map=func_map) + return self._perform_function_sync(func_map=func_map) except Exception as e: logger.error(f"RevEng.AI: Exception during function sync: {e}") - data = GenericApiReturn( - success=False, error_message=f"Exception during function sync: {e}" + return ( + GenericApiReturn(success=False, error_message=f"Exception during function sync: {e}"), + [], + [], ) - return data + def _apply_canonical_names( + self, + needs_canonical: list[tuple[int, int | None, str]], + name_pushbacks: list[RenameInput], + ) -> int: + if not needs_canonical: + return 0 + + mapping: dict[str, str] = self.rename_service.canonicalize_names( + [name for _, _, name in needs_canonical] + ) + + canonicalized: int = 0 + for ea, fid, original in needs_canonical: + canonical: str = mapping.get(original, original) + final: str | None = self.apply_deduped_name(ea, canonical) + if final: + canonicalized += 1 + self.tag_function_as_renamed(final) + if fid is not None: + name_pushbacks.append(RenameInput(ea=ea, new_name=final, function_id=fid)) + return canonicalized def _sync_analysis_data(self, _: threading.Event, func_map: FunctionMapping, on_complete_callback: Callable[[GenericApiReturn[MatchedFunctionSummary]], None]) -> None: - """ - Syncs the analysis data until completion or failure. - """ - response: GenericApiReturn[MatchedFunctionSummary] = self._match_functions(func_map) + response, name_pushbacks, needs_canonical = self._match_functions(func_map) if response.success is False: logger.error(f"failed to sync analysis data due to {response.error_message}") on_complete_callback(response) return - data_types_error: str | None = self.data_types_service.import_data_types( - {int(k): v for k, v in func_map.function_map.items()} - ) - if data_types_error and response.data is not None: - response.data.data_types_error = data_types_error + canonicalized: int = self._apply_canonical_names(needs_canonical, name_pushbacks) + if response.data is not None: + response.data.canonicalized_name_count = canonicalized + + if name_pushbacks: + push_response = self.rename_service.push_remote_names(name_pushbacks) + if response.data is not None and getattr(push_response, "status", False): + response.data.pushed_name_count = len(name_pushbacks) + + matches: dict[int, int] = {int(k): v for k, v in func_map.function_map.items()} + dt_result = self.data_types_service.import_data_types(matches) + if response.data is not None: + response.data.data_types_error = dt_result.error + + surviving: set[int] = { + int(fid) + for fid, vaddr in func_map.function_map.items() + if str(vaddr) in func_map.name_map + } + push_ids: set[int] = (dt_result.remote_absent_ids | dt_result.apply_failed_ids) & surviving + targets: dict[int, int] = {fid: matches[fid] for fid in push_ids if fid in matches} + if targets: + analysis_id: int | None = self.netstore_service.get_analysis_id() + response.data.pushed_type_count = self.variable_sync_service.push_local_function_types_batch( + targets, analysis_id + ) on_complete_callback(response) \ No newline at end of file diff --git a/reai_toolkit/app/services/analysis_sync/schema.py b/reai_toolkit/app/services/analysis_sync/schema.py index 4d332fb..96094b6 100644 --- a/reai_toolkit/app/services/analysis_sync/schema.py +++ b/reai_toolkit/app/services/analysis_sync/schema.py @@ -6,3 +6,7 @@ class MatchedFunctionSummary(BaseModel): unmatched_function_count: int total_function_count: int data_types_error: str | None = None + canonicalized_name_count: int = 0 + deduped_name_count: int = 0 + pushed_name_count: int = 0 + pushed_type_count: int = 0 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 bfcf8af..022924c 100644 --- a/reai_toolkit/app/services/data_types/data_types_service.py +++ b/reai_toolkit/app/services/data_types/data_types_service.py @@ -1,3 +1,5 @@ +from dataclasses import dataclass, field + from revengai import ApiException, Configuration from revengai.exceptions import NotFoundException @@ -17,13 +19,20 @@ FUNCTION_IDS_BATCH_SIZE = 50 +@dataclass +class DataTypesImportResult: + error: str | None = None + remote_absent_ids: set[int] = field(default_factory=set) + apply_failed_ids: set[int] = field(default_factory=set) + + class ImportDataTypesService(IThreadService): def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration) -> None: super().__init__(netstore_service=netstore_service, sdk_config=sdk_config) - def import_data_types(self, matches: dict[int, int]) -> str | None: + def import_data_types(self, matches: dict[int, int]) -> DataTypesImportResult: if len(matches) == 0: - return None + return DataTypesImportResult() idt: ImportDataTypes = ImportDataTypes() matched_function_ids: list[int] = list(matches.keys()) @@ -32,17 +41,29 @@ def import_data_types(self, matches: dict[int, int]) -> str | None: response: FunctionDataTypesList | 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 None + return DataTypesImportResult(remote_absent_ids=set(matched_function_ids)) except ApiException as e: logger.error(f"RevEng.AI: failed to sync function data types: HTTP {e.status} {e.reason}") - return f"Failed to sync function data types: HTTP {e.status} {e.reason}" + return DataTypesImportResult(error=f"Failed to sync function data types: HTTP {e.status} {e.reason}") except Exception as e: logger.error(f"RevEng.AI: failed to sync function data types: {e}") - return f"Failed to sync function data types: {e}" + 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} + if response + else set() + ) + remote_absent_ids: set[int] = set(matched_function_ids) - present_ids + apply_failed_ids: set[int] = set() if response: - idt.execute(response, matched_function_mapping=matches) - return None + apply_failed_ids = idt.execute(response, matched_function_mapping=matches) or set() + + return DataTypesImportResult( + remote_absent_ids=remote_absent_ids, + apply_failed_ids=apply_failed_ids, + ) def _get_data_types(self, function_ids: list[int] | None = None) -> FunctionDataTypesList | None: if not function_ids: diff --git a/reai_toolkit/app/services/rename/rename_service.py b/reai_toolkit/app/services/rename/rename_service.py index b707002..afa7712 100644 --- a/reai_toolkit/app/services/rename/rename_service.py +++ b/reai_toolkit/app/services/rename/rename_service.py @@ -4,12 +4,15 @@ from typing import List, Optional from libbs.decompilers.ida.compat import execute_write +from loguru import logger from revengai import ( BaseResponse, + CanonicalizeNamesInputBody, Configuration, FunctionMapping, FunctionRenameMap, + FunctionsCoreApi, FunctionsListRename, FunctionsRenamingHistoryApi, ) @@ -27,6 +30,7 @@ class RenameService(IThreadService): _rename_last_ts: dict[int, float] = {} _rename_debounce_ms: int = 300 # ignore bursts within 300ms per ea _rename_max_retries: int = 5 + _canonicalize_batch_size: int = 25 def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration): super().__init__(netstore_service=netstore_service, sdk_config=sdk_config) @@ -157,3 +161,27 @@ def _rename_remote_function(self, function_list: list[RenameInput]) -> BaseRespo functions=function_rename_list ) ) + + def push_remote_names(self, renames: list[RenameInput]) -> BaseResponse: + return self._rename_remote_function(renames) + + def canonicalize_names(self, names: list[str]) -> dict[str, str]: + mapping: dict[str, str] = {} + unique: list[str] = [n for n in dict.fromkeys(names) if n] + + for start in range(0, len(unique), self._canonicalize_batch_size): + chunk: list[str] = unique[start:start + self._canonicalize_batch_size] + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + client = FunctionsCoreApi(api_client) + out = client.v3_canonicalize_function_names( + canonicalize_names_input_body=CanonicalizeNamesInputBody(names=chunk) + ) + except Exception as e: + logger.error(f"RevEng.AI: failed to canonicalize names: {e}") + continue + + for result in (out.results or []): + mapping[result.name] = result.canonical_name + + return mapping 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 cc53835..6287951 100644 --- a/reai_toolkit/app/services/variable_sync/variable_sync_service.py +++ b/reai_toolkit/app/services/variable_sync/variable_sync_service.py @@ -53,6 +53,7 @@ class VariableSyncService(IThreadService): _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): super().__init__(netstore_service=netstore_service, sdk_config=sdk_config) @@ -319,3 +320,85 @@ def _patch_header(ft: object, fheader: FunctionHeader) -> None: if arg.type is not None: entry.type = arg.type break + + def push_local_function_types_batch(self, targets: dict[int, int], analysis_id: int | None) -> int: + if not targets or analysis_id is None: + return 0 + + if self._deci is None: + try: + from libbs.api import DecompilerInterface + + self.attach_decompiler(DecompilerInterface.discover(force_decompiler="ida")) + except Exception as e: + logger.error(f"RevEng.AI: could not attach decompiler for type push-back: {e}") + return 0 + + base: int = self._deci.binary_base_addr + infos: dict[int, FunctionInfo] = {} + for function_id, ea in targets.items(): + info = self._build_function_info(ea - base) + 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 + + 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()))) + + 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/reai_toolkit/app/transformations/import_data_types.py b/reai_toolkit/app/transformations/import_data_types.py index 64664d4..0918428 100644 --- a/reai_toolkit/app/transformations/import_data_types.py +++ b/reai_toolkit/app/transformations/import_data_types.py @@ -2,7 +2,7 @@ import libbs.artifacts from libbs.api import DecompilerInterface -from libbs.decompilers.ida.compat import execute_ui +from libbs.decompilers.ida.compat import execute_ui, convert_type_str_to_ida_type from loguru import logger from revengai import ( FunctionArgument, @@ -32,9 +32,10 @@ def __init__(self) -> None: self.deci: DecompilerInterface @execute_ui - def execute(self, functions: FunctionDataTypesList, matched_function_mapping: dict[int, int] = {}) -> None: + def execute(self, functions: FunctionDataTypesList, matched_function_mapping: dict[int, int] = {}) -> set[int]: self.deci = DecompilerInterface.discover(force_decompiler="ida") # type: ignore lookup: dict[str, TaggedDependency] = {} + failed: set[int] = set() for function in functions.items: data_types: FunctionInfo | None = function.data_types @@ -70,18 +71,21 @@ def execute(self, functions: FunctionDataTypesList, matched_function_mapping: di func: FunctionType | None = data_types.func_types if func: - # If we obtained data types from a matched function, we need to make sure we map it to the original effective address. - if matched_function_mapping: - ea: int = matched_function_mapping[function.function_id] - else: - ea: int = func.addr - try: - self.update_function(func, ea) + if matched_function_mapping: + ea: int = matched_function_mapping[function.function_id] + else: + ea: int = func.addr + + if not self.update_function(func, ea) or self._function_types_unparseable(func): + failed.add(function.function_id) except Exception as e: logger.warning( - f"RevEng.AI: skipped data types for {func.name} at ea=0x{ea:x}: {e!r}" + f"RevEng.AI: skipped data types for function {function.function_id}: {e!r}" ) + failed.add(function.function_id) + + return failed def process_dependency( @@ -130,14 +134,14 @@ def update_typedef(self, imported_typedef: TypeDefinition, lookup: dict[str, Tag name=imported_typedef.name, type_=normalized_type ) - def update_function(self, func: FunctionType, ea: int) -> None: + def update_function(self, func: FunctionType, ea: int) -> bool: base_address: int = self.deci.binary_base_addr rva: int = ea - base_address target_func: libbs.artifacts.Function | None = self.deci.functions.get(rva) # type: ignore if target_func is None: logger.warning(f"failed to update function: {func.name} at rva: 0x{rva:0x}") - return + return False target_func.name = func.name target_func.size = func.size @@ -148,6 +152,16 @@ def update_function(self, func: FunctionType, ea: int) -> None: self.update_header(func.header, target_func) self.deci.functions[rva] = target_func + return True + + def _function_types_unparseable(self, func: FunctionType) -> bool: + if not (func.header and func.header.args): + return False + for arg in func.header.args.values(): + type_str: str = self.normalise_type(arg.type) if arg.type else "" + if type_str and convert_type_str_to_ida_type(type_str) is None: + return True + return False def update_header( self, imported_header: FunctionHeader, target_function: libbs.artifacts.Function diff --git a/tests/idalib/analysis_sync/test_sync.py b/tests/idalib/analysis_sync/test_sync.py index 911c322..08e84ee 100644 --- a/tests/idalib/analysis_sync/test_sync.py +++ b/tests/idalib/analysis_sync/test_sync.py @@ -4,6 +4,7 @@ import ida_name import idaapi +import idautils import idc from revengai.models.function_mapping import FunctionMapping @@ -16,6 +17,8 @@ def service(): return AnalysisSyncService( data_types_service=MagicMock(), + rename_service=MagicMock(), + variable_sync_service=MagicMock(), netstore_service=MagicMock(), sdk_config=MagicMock(), ) @@ -30,6 +33,17 @@ def auto_func(loaded_binary): ida_name.set_name(ea, original, ida_name.SN_CHECK | ida_name.SN_AUTO) +@pytest.fixture +def two_funcs(loaded_binary): + eas = list(idautils.Functions()) + assert len(eas) >= 2 + a, b = eas[0], eas[1] + originals = (ida_name.get_name(a), ida_name.get_name(b)) + yield a, b + ida_name.set_name(a, originals[0], ida_name.SN_CHECK | ida_name.SN_AUTO) + ida_name.set_name(b, originals[1], ida_name.SN_CHECK | ida_name.SN_AUTO) + + def test_perform_function_sync_renames_matched_functions(service, auto_func): func_map = FunctionMapping( function_map={}, @@ -37,11 +51,13 @@ def test_perform_function_sync_renames_matched_functions(service, auto_func): name_map={str(auto_func): "synced_fn"}, ) - result = service._perform_function_sync(func_map) + result, pushbacks, needs_canonical = service._perform_function_sync(func_map) assert result.success is True assert result.data.matched_function_count >= 1 assert ida_name.get_name(auto_func) == "synced_fn" + assert pushbacks == [] + assert needs_canonical == [] def test_perform_function_sync_no_matches(service, loaded_binary): @@ -49,8 +65,50 @@ def test_perform_function_sync_no_matches(service, loaded_binary): function_map={}, inverse_function_map={}, name_map={} ) - result = service._perform_function_sync(func_map) + result, pushbacks, needs_canonical = service._perform_function_sync(func_map) assert result.success is True assert result.data.matched_function_count == 0 assert result.data.total_function_count >= 1 + + +def test_perform_function_sync_classifies_invalid_name(service, auto_func): + func_map = FunctionMapping( + function_map={}, + inverse_function_map={str(auto_func): 7}, + name_map={str(auto_func): "bad::name!!"}, + ) + + result, pushbacks, needs_canonical = service._perform_function_sync(func_map) + + assert (auto_func, 7, "bad::name!!") in needs_canonical + assert pushbacks == [] + + +def test_perform_function_sync_dedupes_duplicate_name(service, two_funcs): + ea_a, ea_b = two_funcs + assert ida_name.set_name(ea_a, "dupname", ida_name.SN_CHECK | ida_name.SN_AUTO) + + func_map = FunctionMapping( + function_map={}, + inverse_function_map={str(ea_b): 9}, + name_map={str(ea_b): "dupname"}, + ) + + result, pushbacks, needs_canonical = service._perform_function_sync(func_map) + + assert result.data.deduped_name_count >= 1 + assert ida_name.get_name(ea_b) == "dupname_1" + assert [(p.function_id, p.new_name) for p in pushbacks] == [(9, "dupname_1")] + assert needs_canonical == [] + + +def test_apply_canonical_names_applies_and_collects_pushback(service, auto_func): + service.rename_service.canonicalize_names.return_value = {"bad::name!!": "cleanname"} + pushbacks = [] + + count = service._apply_canonical_names([(auto_func, 7, "bad::name!!")], pushbacks) + + assert count == 1 + assert ida_name.get_name(auto_func) == "cleanname" + assert [(p.function_id, p.new_name) for p in pushbacks] == [(7, "cleanname")] diff --git a/tests/unit/analysis_sync/test_service.py b/tests/unit/analysis_sync/test_service.py index e57f20c..8ce4e09 100644 --- a/tests/unit/analysis_sync/test_service.py +++ b/tests/unit/analysis_sync/test_service.py @@ -5,6 +5,8 @@ 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 +from reai_toolkit.app.services.rename.schema import RenameInput from reai_toolkit.app.core.shared_schema import GenericApiReturn @@ -18,14 +20,31 @@ def netstore(): @pytest.fixture def data_types_service(): dts = MagicMock() - dts.import_data_types.return_value = None + dts.import_data_types.return_value = DataTypesImportResult() return dts @pytest.fixture -def service(netstore, data_types_service): +def rename_service(): + rs = MagicMock() + rs.canonicalize_names.return_value = {} + rs.push_remote_names.return_value = MagicMock(status=True) + return rs + + +@pytest.fixture +def variable_sync_service(): + vs = MagicMock() + vs.push_local_function_types_batch.return_value = 0 + return vs + + +@pytest.fixture +def service(netstore, data_types_service, rename_service, variable_sync_service): return AnalysisSyncService( data_types_service=data_types_service, + rename_service=rename_service, + variable_sync_service=variable_sync_service, netstore_service=netstore, sdk_config=MagicMock(), ) @@ -39,16 +58,20 @@ def _summary() -> MatchedFunctionSummary: def _func_map() -> FunctionMapping: return FunctionMapping.model_construct( - function_map={"1": 0x1000}, inverse_function_map={"4096": 1}, name_map={} + function_map={"1": 0x1000}, inverse_function_map={"4096": 1}, name_map={"4096": "fn"} ) -def test_sync_success_imports_data_types_and_dispatches(service, data_types_service, mocker): - mocker.patch.object( - service, - "_match_functions", - return_value=GenericApiReturn(success=True, data=_summary()), +def _matched(name_pushbacks=None, needs_canonical=None): + return ( + GenericApiReturn(success=True, data=_summary()), + name_pushbacks or [], + needs_canonical or [], ) + + +def test_sync_success_imports_data_types_and_dispatches(service, data_types_service, mocker): + mocker.patch.object(service, "_match_functions", return_value=_matched()) cb = MagicMock() service._sync_analysis_data(MagicMock(), _func_map(), cb) @@ -61,12 +84,10 @@ def test_sync_success_imports_data_types_and_dispatches(service, data_types_serv def test_sync_propagates_data_types_error(service, data_types_service, mocker): - data_types_service.import_data_types.return_value = "403: forbidden" - mocker.patch.object( - service, - "_match_functions", - return_value=GenericApiReturn(success=True, data=_summary()), + data_types_service.import_data_types.return_value = DataTypesImportResult( + error="403: forbidden" ) + mocker.patch.object(service, "_match_functions", return_value=_matched()) cb = MagicMock() service._sync_analysis_data(MagicMock(), _func_map(), cb) @@ -78,7 +99,7 @@ def test_sync_failure_skips_data_types_import(service, data_types_service, mocke mocker.patch.object( service, "_match_functions", - return_value=GenericApiReturn(success=False, error_message="boom"), + return_value=(GenericApiReturn(success=False, error_message="boom"), [], []), ) cb = MagicMock() @@ -93,10 +114,78 @@ def test_match_functions_wraps_exception_as_failure(service, mocker): service, "_perform_function_sync", side_effect=RuntimeError("kaboom") ) - result = service._match_functions(_func_map()) + result, pushbacks, needs_canonical = service._match_functions(_func_map()) assert result.success is False assert "kaboom" in result.error_message + assert pushbacks == [] + assert needs_canonical == [] + + +def test_sync_pushes_deduped_names(service, rename_service, mocker): + pushbacks = [RenameInput(ea=0x1000, new_name="fn_1", function_id=1)] + mocker.patch.object( + service, "_match_functions", return_value=_matched(name_pushbacks=pushbacks) + ) + cb = MagicMock() + + service._sync_analysis_data(MagicMock(), _func_map(), cb) + + rename_service.push_remote_names.assert_called_once_with(pushbacks) + assert cb.call_args[0][0].data.pushed_name_count == 1 + + +def test_sync_canonicalizes_invalid_names(service, rename_service, mocker): + needs = [(0x1000, 1, "bad::name")] + mocker.patch.object( + service, "_match_functions", return_value=_matched(needs_canonical=needs) + ) + mocker.patch.object(service, "apply_deduped_name", return_value="badname") + mocker.patch.object(service, "tag_function_as_renamed") + rename_service.canonicalize_names.return_value = {"bad::name": "badname"} + cb = MagicMock() + + service._sync_analysis_data(MagicMock(), _func_map(), cb) + + rename_service.canonicalize_names.assert_called_once_with(["bad::name"]) + pushed = rename_service.push_remote_names.call_args[0][0] + assert [r.new_name for r in pushed] == ["badname"] + assert cb.call_args[0][0].data.canonicalized_name_count == 1 + + +def test_sync_pushes_local_types_for_absent_and_failed( + service, data_types_service, variable_sync_service, mocker +): + data_types_service.import_data_types.return_value = DataTypesImportResult( + remote_absent_ids={1}, apply_failed_ids=set() + ) + mocker.patch.object(service, "_match_functions", return_value=_matched()) + cb = MagicMock() + + service._sync_analysis_data(MagicMock(), _func_map(), cb) + + variable_sync_service.push_local_function_types_batch.assert_called_once() + targets, analysis_id = variable_sync_service.push_local_function_types_batch.call_args[0] + assert targets == {1: 0x1000} + assert analysis_id == 1234 + assert cb.call_args[0][0].data.pushed_type_count == 0 + + +def test_sync_skips_type_pushback_for_deselected_functions( + service, data_types_service, variable_sync_service, mocker +): + func_map = FunctionMapping.model_construct( + function_map={"1": 0x1000}, inverse_function_map={"4096": 1}, name_map={} + ) + data_types_service.import_data_types.return_value = DataTypesImportResult( + remote_absent_ids={1} + ) + mocker.patch.object(service, "_match_functions", return_value=_matched()) + cb = MagicMock() + + service._sync_analysis_data(MagicMock(), func_map, cb) + + variable_sync_service.push_local_function_types_batch.assert_not_called() def test_get_function_matches_invokes_callback_with_map(service, mocker): diff --git a/tests/unit/data_types/test_service.py b/tests/unit/data_types/test_service.py index ebd6e7b..f2b03b7 100644 --- a/tests/unit/data_types/test_service.py +++ b/tests/unit/data_types/test_service.py @@ -1,7 +1,6 @@ from unittest.mock import MagicMock import pytest -from revengai import ApiException 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 @@ -99,42 +98,71 @@ def test_skips_unsuccessful_chunks(service, sdk): assert [item.function_id for item in result.items] == [1, 2, 1, 2] -def test_import_data_types_returns_none_on_success(service, sdk, mocker): +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)]) - apply = mocker.patch.object(svc_mod.ImportDataTypes, "execute") + apply = mocker.patch.object(svc_mod.ImportDataTypes, "execute", return_value=set()) + + result = service.import_data_types({1: 0x1000}) - assert service.import_data_types({1: 0x1000}) is None apply.assert_called_once() + assert result.error is None + assert result.remote_absent_ids == {1} + assert result.apply_failed_ids == set() + + +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]) + mocker.patch.object(svc_mod.ImportDataTypes, "execute", return_value={1}) + + result = service.import_data_types({1: 0x1000}) + assert result.remote_absent_ids == set() + assert result.apply_failed_ids == {1} -def test_import_data_types_returns_none_for_empty_matches(service, sdk): - assert service.import_data_types({}) is None + +def test_import_data_types_empty_matches(service, sdk): + result = service.import_data_types({}) + + 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() -def test_import_data_types_returns_single_error_on_forbidden(service, sdk, mocker): +def test_import_data_types_returns_error_on_forbidden(service, sdk, mocker): sdk.list_function_data_types_for_functions.side_effect = ForbiddenException( status=403, reason="Forbidden" ) apply = mocker.patch.object(svc_mod.ImportDataTypes, "execute") matches = {fid: fid * 16 for fid in range(FUNCTION_IDS_BATCH_SIZE * 3)} - err = service.import_data_types(matches) + result = service.import_data_types(matches) - assert err is not None - assert "403" in err - assert "Forbidden" in err + assert result.error is not None + assert "403" in result.error + 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 -def test_import_data_types_treats_not_found_as_warning(service, sdk, mocker): +def test_import_data_types_treats_not_found_as_all_absent(service, sdk, mocker): sdk.list_function_data_types_for_functions.side_effect = NotFoundException( status=404, reason="Not Found" ) apply = mocker.patch.object(svc_mod.ImportDataTypes, "execute") - assert service.import_data_types({1: 0x1000}) is None + result = service.import_data_types({1: 0x1000, 2: 0x2000}) + + assert result.error is None + assert result.remote_absent_ids == {1, 2} apply.assert_not_called() @@ -142,8 +170,8 @@ def test_import_data_types_returns_error_on_unexpected_exception(service, sdk, m sdk.list_function_data_types_for_functions.side_effect = RuntimeError("boom") apply = mocker.patch.object(svc_mod.ImportDataTypes, "execute") - err = service.import_data_types({1: 0x1000}) + result = service.import_data_types({1: 0x1000}) - assert err is not None - assert "boom" in err + assert result.error is not None + assert "boom" in result.error apply.assert_not_called() diff --git a/tests/unit/rename/test_service.py b/tests/unit/rename/test_service.py index c937092..7f8ec7f 100644 --- a/tests/unit/rename/test_service.py +++ b/tests/unit/rename/test_service.py @@ -1,4 +1,5 @@ import queue +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -110,3 +111,53 @@ def test_enqueue_rename_debounces_rapid_duplicates(service, mocker): second = service._rename_q.get_nowait() assert [r.new_name for r in first] == ["a"] assert second == [] + + +def test_push_remote_names_delegates_to_remote_rename(service, ida_calls): + _, remote = ida_calls + renames = [RenameInput(ea=0x10, new_name="foo", function_id=1)] + + resp = service.push_remote_names(renames) + + remote.assert_called_once_with(renames) + assert resp.status is True + + +def test_canonicalize_names_maps_and_chunks_at_25(service, mocker): + mocker.patch.object(RenameService, "yield_api_client") + api_class = mocker.patch( + "reai_toolkit.app.services.rename.rename_service.FunctionsCoreApi" + ) + api = MagicMock() + api_class.return_value = api + + def _canon(canonicalize_names_input_body): + out = MagicMock() + out.results = [ + SimpleNamespace(name=n, canonical_name=n.upper()) + for n in canonicalize_names_input_body.names + ] + return out + + api.v3_canonicalize_function_names.side_effect = _canon + + names = [f"n{i}" for i in range(30)] + mapping = service.canonicalize_names(names) + + assert api.v3_canonicalize_function_names.call_count == 2 + assert mapping["n0"] == "N0" + assert len(mapping) == 30 + + +def test_canonicalize_names_skips_failed_chunk(service, mocker): + mocker.patch.object(RenameService, "yield_api_client") + api_class = mocker.patch( + "reai_toolkit.app.services.rename.rename_service.FunctionsCoreApi" + ) + api = MagicMock() + api_class.return_value = api + api.v3_canonicalize_function_names.side_effect = RuntimeError("boom") + + mapping = service.canonicalize_names(["a", "b"]) + + assert mapping == {} diff --git a/tests/unit/variable_sync/test_service.py b/tests/unit/variable_sync/test_service.py index a898e8d..9c08b0f 100644 --- a/tests/unit/variable_sync/test_service.py +++ b/tests/unit/variable_sync/test_service.py @@ -1,4 +1,5 @@ import queue +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -107,6 +108,82 @@ def test_patch_header_updates_arg_and_return_type(service): assert info.func_types.type == "int" +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 + + +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 + + updated = service.push_local_function_types_batch({1: 0x401000}, analysis_id=9) + + 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 + + +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 + + service.push_local_function_types_batch({1: 0x1000}, analysis_id=9) + + body = sdk.batch_update_function_data_types.call_args.kwargs[ + "batch_update_data_types_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 + + def test_patch_stack_var_type_change_keeps_name(service): info = _function_info(stack_vars={"-0x20": _sdk_stack_var(-32, "local_20", "char")}) diff --git a/uv.lock b/uv.lock index d7e8c3c..8c782da 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.110.0" }, + { name = "revengai", specifier = ">=3.113.0" }, ] [package.metadata.requires-dev] @@ -855,7 +855,7 @@ wheels = [ [[package]] name = "revengai" -version = "3.110.0" +version = "3.113.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/63/e0/8ab3b3452b020d9167ac82f10f4b89c6e81dc4fda74667476028daf8c02b/revengai-3.110.0.tar.gz", hash = "sha256:af728016a48a1a50f9e58ff0ffe83f81c9ef794e2d2e0e7916ee527a291dbeb0", size = 372387, upload-time = "2026-07-08T07:45:35.775Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5c/c1/98891a89e6f19ca9959d6468ca26c22314ef878b25ecfc475ec03ca7db38/revengai-3.113.0.tar.gz", hash = "sha256:29267fc9ec02690895fb1a675b66b6be2bf92b80fc90599592101efb96a9370b", size = 379520, upload-time = "2026-07-09T08:25:44.645Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/20/77/e7341b7b8a649d2747c850fcd7be46a13d366e3fef0347d8a4701747c8b0/revengai-3.110.0-py3-none-any.whl", hash = "sha256:b389cc6039b96f44d8fd4bb96f9e7d74540c4462986ff81a2c171357d2e56587", size = 965938, upload-time = "2026-07-08T07:45:34.518Z" }, + { url = "https://files.pythonhosted.org/packages/99/b7/eac5d9f1b90e400a8a1db164d3343c234ee8e37ca32bbb94c1d534cb489e/revengai-3.113.0-py3-none-any.whl", hash = "sha256:770201b4f210cab09bd988982a957230ee77c7564d904187ebf11ac802a22d37", size = 983788, upload-time = "2026-07-09T08:25:43.027Z" }, ] [[package]] From b4e2297e62d6490c3be731449e372efd8ccbf87d Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Thu, 9 Jul 2026 13:12:34 +0100 Subject: [PATCH 2/3] fix: dt sync bug --- reai_toolkit/app/interfaces/base_service.py | 4 +- .../services/analysis_sync/analysis_sync.py | 15 +++++-- .../variable_sync/variable_sync_service.py | 40 ++++++++----------- tests/idalib/analysis_sync/test_sync.py | 8 ++-- tests/unit/variable_sync/test_service.py | 28 +++++++++++-- 5 files changed, 58 insertions(+), 37 deletions(-) diff --git a/reai_toolkit/app/interfaces/base_service.py b/reai_toolkit/app/interfaces/base_service.py index 23c9609..39d3ce7 100644 --- a/reai_toolkit/app/interfaces/base_service.py +++ b/reai_toolkit/app/interfaces/base_service.py @@ -71,7 +71,7 @@ def update_function_name(ea: int, new_name: str, check_user_flags: bool = False) # SN_CHECK: check for validity # SN_AUTO: mark as auto-generated name # SN_NODUMMY: Prevents warning "can't rename byte as '' because the name has a reserved prefix". - flags: int = ida_name.SN_CHECK | ida_name.SN_AUTO | ida_name.SN_NODUMMY + flags: int = ida_name.SN_CHECK | ida_name.SN_AUTO | ida_name.SN_NODUMMY | ida_name.SN_NOWARN return ida_name.set_name(ea, new_name, flags) @@ -87,7 +87,7 @@ def is_protected_user_name(ea: int) -> bool: @staticmethod @execute_write def apply_deduped_name(ea: int, desired_name: str) -> str | None: - flags: int = ida_name.SN_CHECK | ida_name.SN_AUTO | ida_name.SN_NODUMMY + flags: int = ida_name.SN_CHECK | ida_name.SN_AUTO | ida_name.SN_NODUMMY | ida_name.SN_NOWARN candidate: str = desired_name suffix: int = 1 diff --git a/reai_toolkit/app/services/analysis_sync/analysis_sync.py b/reai_toolkit/app/services/analysis_sync/analysis_sync.py index 571a638..b96df8f 100644 --- a/reai_toolkit/app/services/analysis_sync/analysis_sync.py +++ b/reai_toolkit/app/services/analysis_sync/analysis_sync.py @@ -136,7 +136,9 @@ def _perform_function_sync( func_map: FunctionMapping, ) -> tuple[GenericApiReturn[MatchedFunctionSummary], list[RenameInput], list[tuple[int, int | None, str]]]: local_vaddr_to_matched_name: dict[str, str] = func_map.name_map - inverse_map: dict[str, int] = func_map.inverse_function_map + addr_to_function_id: dict[int, int] = { + int(addr): int(fid) for fid, addr in func_map.function_map.items() + } logger.info(f"RevEng.AI: Retrieved {len(local_vaddr_to_matched_name)} functions from analysis") @@ -157,7 +159,7 @@ def _perform_function_sync( if new_name: matched_function_count += 1 if new_name != old_name: - fid: int | None = inverse_map.get(local_vaddr_str) + fid: int | None = addr_to_function_id.get(local_vaddr) if self.update_function_name(local_vaddr, new_name, check_user_flags=True): self.tag_function_as_renamed(new_name) elif self.is_protected_user_name(local_vaddr): @@ -240,8 +242,13 @@ def _sync_analysis_data(self, _: threading.Event, func_map: FunctionMapping, on_ if name_pushbacks: push_response = self.rename_service.push_remote_names(name_pushbacks) - if response.data is not None and getattr(push_response, "status", False): - response.data.pushed_name_count = len(name_pushbacks) + if getattr(push_response, "status", False): + if response.data is not None: + response.data.pushed_name_count = len(name_pushbacks) + else: + logger.error( + f"RevEng.AI: failed to push {len(name_pushbacks)} corrected name(s) to the platform" + ) matches: dict[int, int] = {int(k): v for k, v in func_map.function_map.items()} dt_result = self.data_types_service.import_data_types(matches) 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 6287951..c439c16 100644 --- a/reai_toolkit/app/services/variable_sync/variable_sync_service.py +++ b/reai_toolkit/app/services/variable_sync/variable_sync_service.py @@ -14,18 +14,14 @@ BatchUpdateDataTypesItem, BatchUpdateDataTypesOutputBody, Configuration, - Enumeration, + FunctionDependency, FunctionInfo, - V2FunctionInfoFuncDepsInner, FunctionsDataTypesApi, FunctionType, - Structure, - TypeDefinition, ) 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.models.structure_member import StructureMember from reai_toolkit.app.core.netstore_service import SimpleNetStore from reai_toolkit.app.interfaces.thread_service import IThreadService @@ -195,7 +191,7 @@ def _collect_func_deps(self, func_type: FunctionType) -> list: 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()) - deps: dict[str, V2FunctionInfoFuncDepsInner] = {} + deps: dict[str, FunctionDependency] = {} seen: set[str] = set() while pending and len(deps) < 200: name = self._base_type_name(pending.pop()) @@ -210,38 +206,32 @@ def _collect_func_deps(self, func_type: FunctionType) -> list: return list(deps.values()) - def _resolve_type(self, name: str) -> Tuple[Optional[V2FunctionInfoFuncDepsInner], list]: + def _resolve_type(self, name: str) -> Tuple[Optional[FunctionDependency], list]: artifact = _read_named_type(self._deci, name) if isinstance(artifact, Typedef): return ( - V2FunctionInfoFuncDepsInner( - TypeDefinition(name=artifact.name, type=artifact.type or "", artifact_type="Typedef") - ), + FunctionDependency(name=artifact.name, type=artifact.type or "", artifact_type="Typedef"), [artifact.type], ) if isinstance(artifact, Struct): members = { - hex(member.offset): StructureMember( - name=member.name or "", - offset=member.offset, - type=member.type or "", - size=member.size or 0, - ) + 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 ( - V2FunctionInfoFuncDepsInner( - Structure(name=artifact.name, size=artifact.size, members=members, artifact_type="Struct") - ), + FunctionDependency(name=artifact.name, size=artifact.size, members=members, artifact_type="Struct"), referenced, ) if isinstance(artifact, Enum): members = {str(key): int(value) for key, value in (artifact.members or {}).items()} return ( - V2FunctionInfoFuncDepsInner( - Enumeration(name=artifact.name, members=members, artifact_type="Enum") - ), + FunctionDependency(name=artifact.name, members=members, artifact_type="Enum"), [], ) return None, [] @@ -337,7 +327,11 @@ def push_local_function_types_batch(self, targets: dict[int, int], analysis_id: base: int = self._deci.binary_base_addr infos: dict[int, FunctionInfo] = {} for function_id, ea in targets.items(): - info = self._build_function_info(ea - base) + try: + info = self._build_function_info(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 diff --git a/tests/idalib/analysis_sync/test_sync.py b/tests/idalib/analysis_sync/test_sync.py index 08e84ee..0e6f07f 100644 --- a/tests/idalib/analysis_sync/test_sync.py +++ b/tests/idalib/analysis_sync/test_sync.py @@ -74,8 +74,8 @@ def test_perform_function_sync_no_matches(service, loaded_binary): def test_perform_function_sync_classifies_invalid_name(service, auto_func): func_map = FunctionMapping( - function_map={}, - inverse_function_map={str(auto_func): 7}, + function_map={"7": auto_func}, + inverse_function_map={}, name_map={str(auto_func): "bad::name!!"}, ) @@ -90,8 +90,8 @@ def test_perform_function_sync_dedupes_duplicate_name(service, two_funcs): assert ida_name.set_name(ea_a, "dupname", ida_name.SN_CHECK | ida_name.SN_AUTO) func_map = FunctionMapping( - function_map={}, - inverse_function_map={str(ea_b): 9}, + function_map={"9": ea_b}, + inverse_function_map={}, name_map={str(ea_b): "dupname"}, ) diff --git a/tests/unit/variable_sync/test_service.py b/tests/unit/variable_sync/test_service.py index 9c08b0f..4621024 100644 --- a/tests/unit/variable_sync/test_service.py +++ b/tests/unit/variable_sync/test_service.py @@ -184,6 +184,26 @@ def test_push_local_types_batch_retries_on_version_conflict(service, sdk, mocker assert sdk.batch_update_function_data_types.call_count == 2 +def test_resolve_type_builds_valid_function_info_dependency(service, mocker): + from libbs.artifacts import Typedef + from revengai import FunctionDependency, FunctionInfo + + 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", + } + + def test_patch_stack_var_type_change_keeps_name(service): info = _function_info(stack_vars={"-0x20": _sdk_stack_var(-32, "local_20", "char")}) @@ -346,8 +366,8 @@ def test_collect_func_deps_resolves_typedef_chain(service, mocker): deps = service._collect_func_deps(func_type) - assert sorted(d.actual_instance.name for d in deps) == ["__dev_t", "dev_t"] - assert all(d.actual_instance.to_dict()["artifact_type"] == "Typedef" for d in deps) + assert sorted(d.name for d in deps) == ["__dev_t", "dev_t"] + assert all(d.artifact_type == "Typedef" for d in deps) def test_collect_func_deps_resolves_struct_members(service, mocker): @@ -366,9 +386,9 @@ def test_collect_func_deps_resolves_struct_members(service, mocker): deps = service._collect_func_deps(func_type) - by_name = {d.actual_instance.name: d.actual_instance for d in deps} + by_name = {d.name: d for d in deps} assert set(by_name) == {"mystruct", "myint"} - assert by_name["mystruct"].members["0x0"].type == "myint" + assert by_name["mystruct"].members["0x0"]["type"] == "myint" def test_push_change_no_analysis_id_does_nothing(service, sdk, netstore): From 7858eea23a19f061667f161ea93c1826ea148c68 Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Thu, 9 Jul 2026 14:00:00 +0100 Subject: [PATCH 3/3] fix: func rename failures --- .../app/services/analysis_sync/analysis_sync.py | 14 +++++++------- .../app/services/rename/rename_service.py | 14 ++++++++++++-- tests/unit/analysis_sync/test_service.py | 2 +- tests/unit/rename/test_service.py | 17 +++++++++++++++-- 4 files changed, 35 insertions(+), 12 deletions(-) diff --git a/reai_toolkit/app/services/analysis_sync/analysis_sync.py b/reai_toolkit/app/services/analysis_sync/analysis_sync.py index b96df8f..92248d8 100644 --- a/reai_toolkit/app/services/analysis_sync/analysis_sync.py +++ b/reai_toolkit/app/services/analysis_sync/analysis_sync.py @@ -241,13 +241,13 @@ def _sync_analysis_data(self, _: threading.Event, func_map: FunctionMapping, on_ response.data.canonicalized_name_count = canonicalized if name_pushbacks: - push_response = self.rename_service.push_remote_names(name_pushbacks) - if getattr(push_response, "status", False): - if response.data is not None: - response.data.pushed_name_count = len(name_pushbacks) - else: - logger.error( - f"RevEng.AI: failed to push {len(name_pushbacks)} corrected name(s) to the platform" + pushed: int = self.rename_service.push_remote_names(name_pushbacks) + if response.data is not None: + response.data.pushed_name_count = pushed + if pushed < len(name_pushbacks): + logger.warning( + f"RevEng.AI: pushed {pushed}/{len(name_pushbacks)} corrected name(s); " + f"{len(name_pushbacks) - pushed} rejected (functions not editable on the platform)" ) matches: dict[int, int] = {int(k): v for k, v in func_map.function_map.items()} diff --git a/reai_toolkit/app/services/rename/rename_service.py b/reai_toolkit/app/services/rename/rename_service.py index afa7712..4bed6b1 100644 --- a/reai_toolkit/app/services/rename/rename_service.py +++ b/reai_toolkit/app/services/rename/rename_service.py @@ -31,6 +31,7 @@ class RenameService(IThreadService): _rename_debounce_ms: int = 300 # ignore bursts within 300ms per ea _rename_max_retries: int = 5 _canonicalize_batch_size: int = 25 + _rename_batch_size: int = 50 def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration): super().__init__(netstore_service=netstore_service, sdk_config=sdk_config) @@ -162,8 +163,17 @@ def _rename_remote_function(self, function_list: list[RenameInput]) -> BaseRespo ) ) - def push_remote_names(self, renames: list[RenameInput]) -> BaseResponse: - return self._rename_remote_function(renames) + 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) + continue + for rename in chunk: + if getattr(self._rename_remote_function([rename]), "status", False): + pushed += 1 + return pushed def canonicalize_names(self, names: list[str]) -> dict[str, str]: mapping: dict[str, str] = {} diff --git a/tests/unit/analysis_sync/test_service.py b/tests/unit/analysis_sync/test_service.py index 8ce4e09..6af98ad 100644 --- a/tests/unit/analysis_sync/test_service.py +++ b/tests/unit/analysis_sync/test_service.py @@ -28,7 +28,7 @@ def data_types_service(): def rename_service(): rs = MagicMock() rs.canonicalize_names.return_value = {} - rs.push_remote_names.return_value = MagicMock(status=True) + rs.push_remote_names.side_effect = lambda renames: len(renames) return rs diff --git a/tests/unit/rename/test_service.py b/tests/unit/rename/test_service.py index 7f8ec7f..529ea83 100644 --- a/tests/unit/rename/test_service.py +++ b/tests/unit/rename/test_service.py @@ -117,10 +117,23 @@ def test_push_remote_names_delegates_to_remote_rename(service, ida_calls): _, remote = ida_calls renames = [RenameInput(ea=0x10, new_name="foo", function_id=1)] - resp = service.push_remote_names(renames) + pushed = service.push_remote_names(renames) + assert pushed == 1 remote.assert_called_once_with(renames) - assert resp.status is True + + +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) + remote.side_effect = [bad, ok, bad, ok] + + pushed = service.push_remote_names(renames) + + assert pushed == 2 + assert remote.call_count == 4 def test_canonicalize_names_maps_and_chunks_at_25(service, mocker):