feat(PLU-314): syanc analysis updates

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",
 ]
 

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(

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()
 

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
@@ -70,10 +71,39 @@ 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 '<func_name>' 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)
 
+    @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 | ida_name.SN_NOWARN
+
+        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
     # =========================================

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,138 @@ 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
+        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")
 
-        # 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 = 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):
+                        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:
+            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()}
+        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)

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

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: