feat(PLU-301): auto unstrip re-sync

reai_toolkit/app/app.py

@@ -3,6 +3,9 @@
 from reai_toolkit.app.services.analysis_status.analysis_status import (
     AnalysisStatusService,
 )
+from reai_toolkit.app.services.auto_unstrip_status.auto_unstrip_status import (
+    AutoUnstripStatusService,
+)
 from reai_toolkit.app.services.analysis_sync.analysis_sync import AnalysisSyncService
 from reai_toolkit.app.services.auth.auth_service import AuthService
 from reai_toolkit.app.services.chat.chat_service import ChatService
@@ -41,6 +44,9 @@ def __init__(self, ida_version: str = "UNKNOWN", plugin_version: str = "UNKNOWN"
         self.analysis_status_service = AnalysisStatusService(
             netstore_service=self.netstore_service, sdk_config=sdk_config
         )
+        self.auto_unstrip_status_service = AutoUnstripStatusService(
+            netstore_service=self.netstore_service, sdk_config=sdk_config
+        )
         self.data_types_service = ImportDataTypesService(
             netstore_service=self.netstore_service, sdk_config=sdk_config
         )

reai_toolkit/app/coordinator.py

@@ -5,6 +5,9 @@
 from reai_toolkit.app.coordinators.about_coordinator import AboutCoordinator
 from reai_toolkit.app.coordinators.ai_decomp_coordinator import AiDecompCoordinator
 from reai_toolkit.app.coordinators.auth_coordinator import AuthCoordinator
+from reai_toolkit.app.coordinators.auto_unstrip_coordinator import (
+    AutoUnstripStatusCoordinator,
+)
 from reai_toolkit.app.coordinators.base_coordinator import BaseCoordinator
 from reai_toolkit.app.coordinators.chat_coordinator import ChatCoordinator
 from reai_toolkit.app.coordinators.create_analysis_coordinator import (
@@ -44,12 +47,22 @@ def __init__(self, app, factory, log):
             log=log,
             analysis_sync_service=app.analysis_sync_service,
         )
+        self.auto_unstrip_statusc: AutoUnstripStatusCoordinator = (
+            AutoUnstripStatusCoordinator(
+                app=app,
+                factory=factory,
+                log=log,
+                auto_unstrip_status_service=app.auto_unstrip_status_service,
+                analysis_sync_coord=self.sync_analysisc,
+            )
+        )
         self.poll_statusc: AnalysisStatusCoordinator = AnalysisStatusCoordinator(
             app=app,
             factory=factory,
             log=log,
             analysis_status_service=app.analysis_status_service,
             analysis_sync_coord=self.sync_analysisc,
+            auto_unstrip_coord=self.auto_unstrip_statusc,
         )
         self.create_analysisc: CreateAnalysisCoordinator = CreateAnalysisCoordinator(
             app=app, factory=factory, log=log, analysis_status_coord=self.poll_statusc
@@ -66,6 +79,7 @@ def __init__(self, app, factory, log):
                 log=log,
                 existing_analyses_service=app.existing_analyses_service,
                 analysis_sync_coord=self.sync_analysisc,
+                auto_unstrip_coord=self.auto_unstrip_statusc,
             )
         )
 

reai_toolkit/app/coordinators/auto_unstrip_coordinator.py

@@ -0,0 +1,69 @@
+from typing import TYPE_CHECKING
+
+from reai_toolkit.app.coordinators.base_coordinator import BaseCoordinator
+from reai_toolkit.app.coordinators.sync_analysis_coordinator import (
+    AnalysisSyncCoordinator,
+)
+from reai_toolkit.app.core.shared_schema import GenericApiReturn
+from reai_toolkit.app.services.auto_unstrip_status.auto_unstrip_status import (
+    AutoUnstripStatusService,
+)
+
+if TYPE_CHECKING:
+    from reai_toolkit.app.app import App
+    from reai_toolkit.app.factory import DialogFactory
+
+
+class AutoUnstripStatusCoordinator(BaseCoordinator):
+    auto_unstrip_status_service: AutoUnstripStatusService
+    analysis_sync_coord: AnalysisSyncCoordinator
+
+    def __init__(
+        self,
+        *,
+        app: "App",
+        factory: "DialogFactory",
+        log,
+        auto_unstrip_status_service: AutoUnstripStatusService,
+        analysis_sync_coord: AnalysisSyncCoordinator,
+    ):
+        super().__init__(app=app, factory=factory, log=log)
+
+        self.auto_unstrip_status_service = auto_unstrip_status_service
+        self.analysis_sync_coord = analysis_sync_coord
+
+    def run_dialog(self) -> None:
+        pass
+
+    def is_authed(self) -> bool:
+        return self.app.auth_service.is_authenticated()
+
+    def is_active_worker(self) -> bool:
+        return self.auto_unstrip_status_service.is_worker_running()
+
+    def poll_and_resync(
+        self, analysis_id: int, attach_to_existing_analysis: bool = False
+    ) -> None:
+        self.auto_unstrip_status_service.start_polling(
+            analysis_id=analysis_id,
+            thread_callback=lambda generic_return: self._on_complete(
+                generic_return, attach_to_existing_analysis
+            ),
+            resync_if_already_complete=not attach_to_existing_analysis,
+        )
+
+    def _on_complete(
+        self, generic_return: GenericApiReturn[int], attach_to_existing_analysis: bool
+    ) -> None:
+        if generic_return.success is False:
+            self.log.warning(
+                generic_return.error_message or "failed to poll auto-unstrip status"
+            )
+            return
+
+        self.show_info_dialog(
+            msg="Auto-unstrip has finished. Re-syncing recovered function names and data types."
+        )
+        self.analysis_sync_coord.sync_analysis(
+            attach_to_existing_analysis=attach_to_existing_analysis
+        )

reai_toolkit/app/coordinators/detach_coordinator.py

@@ -21,6 +21,7 @@ def __init__(
     def _detach_analysis(self) -> None:
         self.app.analysis_sync_service.stop_worker()
         self.app.analysis_status_service.stop_worker()
+        self.app.auto_unstrip_status_service.stop_worker()
         self.app.ai_decomp_service.stop_worker()
         self.app.rename_service.stop_worker()
         self.app.matching_service.stop_worker()

reai_toolkit/app/coordinators/existing_analysis_coordinator.py

@@ -2,6 +2,9 @@
 
 import idaapi
 
+from reai_toolkit.app.coordinators.auto_unstrip_coordinator import (
+    AutoUnstripStatusCoordinator,
+)
 from reai_toolkit.app.coordinators.base_coordinator import BaseCoordinator
 from reai_toolkit.app.coordinators.sync_analysis_coordinator import (
     AnalysisSyncCoordinator,
@@ -20,6 +23,7 @@
 class ExistingAnalysesCoordinator(BaseCoordinator):
     existing_analyses_service: ExistingAnalysesService
     analysis_sync_coord: AnalysisSyncCoordinator
+    auto_unstrip_coord: AutoUnstripStatusCoordinator
 
     def __init__(
         self,
@@ -29,11 +33,13 @@ def __init__(
         log,
         existing_analyses_service: ExistingAnalysesService,
         analysis_sync_coord: AnalysisSyncCoordinator,
+        auto_unstrip_coord: AutoUnstripStatusCoordinator,
     ):
         super().__init__(app=app, factory=factory, log=log)
 
         self.existing_analyses_service = existing_analyses_service
         self.analysis_sync_coord = analysis_sync_coord
+        self.auto_unstrip_coord = auto_unstrip_coord
 
     def run_dialog(self) -> None:
         response = self.existing_analyses_service.fetch_analyses_same_hash(
@@ -65,6 +71,10 @@ def run_dialog(self) -> None:
 
         self.analysis_sync_coord.sync_analysis(attach_to_existing_analysis=True)
 
+        self.auto_unstrip_coord.poll_and_resync(
+            analysis_id=data.analysis_id, attach_to_existing_analysis=True
+        )
+
 
     def is_authed(self) -> bool:
         return self.app.auth_service.is_authenticated()

reai_toolkit/app/coordinators/poll_status_coordinator.py

@@ -1,5 +1,8 @@
 from typing import TYPE_CHECKING
 
+from reai_toolkit.app.coordinators.auto_unstrip_coordinator import (
+    AutoUnstripStatusCoordinator,
+)
 from reai_toolkit.app.coordinators.base_coordinator import BaseCoordinator
 from reai_toolkit.app.coordinators.sync_analysis_coordinator import (
     AnalysisSyncCoordinator,
@@ -17,6 +20,7 @@
 class AnalysisStatusCoordinator(BaseCoordinator):
     analysis_status_service: AnalysisStatusService
     analysis_sync_coord: AnalysisSyncCoordinator
+    auto_unstrip_coord: AutoUnstripStatusCoordinator
 
     def __init__(
         self,
@@ -26,11 +30,13 @@ def __init__(
         log,
         analysis_status_service: AnalysisStatusService,
         analysis_sync_coord: AnalysisSyncCoordinator,
+        auto_unstrip_coord: AutoUnstripStatusCoordinator,
     ):
         super().__init__(app=app, factory=factory, log=log)
 
         self.analysis_status_service = analysis_status_service
         self.analysis_sync_coord = analysis_sync_coord
+        self.auto_unstrip_coord = auto_unstrip_coord
 
     def run_dialog(self) -> None:
         pass
@@ -59,4 +65,10 @@ def _on_complete(self, generic_return: GenericApiReturn[int]) -> None:
 
         self.analysis_sync_coord.sync_analysis()
 
+        analysis_id: int | None = generic_return.data
+        if analysis_id is not None:
+            self.auto_unstrip_coord.poll_and_resync(
+                analysis_id=analysis_id, attach_to_existing_analysis=False
+            )
+
         self.refresh_disassembly_view()

reai_toolkit/app/coordinators/sync_analysis_coordinator.py

@@ -104,6 +104,11 @@ def _on_complete(
                 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.missing_symbol_name_count:
+                msg += (
+                    f"\n{data.missing_symbol_name_count} function(s) in the remote analysis "
+                    f"with missing symbol names."
+                )
             if data.canonicalized_name_count or data.deduped_name_count:
                 msg += (
                     f"\nCanonicalized {data.canonicalized_name_count} and disambiguated "

reai_toolkit/app/services/analysis_sync/analysis_sync.py

@@ -144,6 +144,7 @@ def _perform_function_sync(
 
         matched_function_count: int = 0
         unmatched_function_count: int = 0
+        missing_symbol_name_count: int = 0
         total_function_count: int = 0
         deduped_name_count: int = 0
 
@@ -177,17 +178,21 @@ def _perform_function_sync(
                                     )
                         else:
                             needs_canonical.append((local_vaddr, fid, new_name))
+            elif local_vaddr_str in local_vaddr_to_matched_name:
+                missing_symbol_name_count += 1
             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")
+        logger.info(f"RevEng.AI: {missing_symbol_name_count} functions with missing symbol names")
 
         summary: MatchedFunctionSummary = MatchedFunctionSummary(
             matched_function_count=matched_function_count,
             unmatched_function_count=unmatched_function_count,
+            missing_symbol_name_count=missing_symbol_name_count,
             total_function_count=total_function_count,
             deduped_name_count=deduped_name_count,
         )

reai_toolkit/app/services/analysis_sync/schema.py

@@ -5,6 +5,7 @@ class MatchedFunctionSummary(BaseModel):
     matched_function_count: int
     unmatched_function_count: int
     total_function_count: int
+    missing_symbol_name_count: int = 0
     data_types_error: str | None = None
     canonicalized_name_count: int = 0
     deduped_name_count: int = 0

reai_toolkit/app/services/auto_unstrip_status/auto_unstrip_status.py

@@ -0,0 +1,84 @@
+import threading
+import time
+from typing import Any, Callable
+
+from loguru import logger
+from revengai import AnalysesCoreApi, AutoUnstripStatusOutputBody, Configuration
+
+from reai_toolkit.app.core.netstore_service import SimpleNetStore
+from reai_toolkit.app.core.shared_schema import GenericApiReturn
+from reai_toolkit.app.interfaces.thread_service import IThreadService
+
+POLL_INTERVAL_SECONDS = 5
+STATUS_COMPLETED = "COMPLETED"
+STATUS_FAILED = "FAILED"
+
+
+class AutoUnstripStatusService(IThreadService):
+    _thread_callback: Callable[..., Any] = None
+
+    def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration):
+        super().__init__(netstore_service=netstore_service, sdk_config=sdk_config)
+
+    def call_callback(self, generic_return: GenericApiReturn) -> None:
+        self._thread_callback(generic_return)
+
+    def start_polling(
+        self,
+        analysis_id: int,
+        thread_callback: Callable[..., Any],
+        resync_if_already_complete: bool,
+    ) -> None:
+        self._thread_callback = thread_callback
+        self.stop_worker()
+        self.start_worker(
+            target=self._poll_auto_unstrip_status,
+            args=(analysis_id, resync_if_already_complete),
+        )
+
+    def _api_get_status(self, analysis_id: int) -> str | None:
+        with self.yield_api_client(sdk_config=self.sdk_config) as api_client:
+            analyses_client = AnalysesCoreApi(api_client)
+            body: AutoUnstripStatusOutputBody = (
+                analyses_client.v3_get_analysis_auto_unstrip_status(analysis_id=analysis_id)
+            )
+            return body.status if body else None
+
+    def _poll_auto_unstrip_status(
+        self,
+        stop_event: threading.Event,
+        analysis_id: int,
+        resync_if_already_complete: bool,
+    ) -> None:
+        first_poll: bool = True
+        while not stop_event.is_set():
+            status_response = self.api_request_returning(
+                fn=lambda: self._api_get_status(analysis_id)
+            )
+
+            if not status_response.success:
+                self.call_callback(generic_return=status_response)
+                return
+
+            status: str | None = status_response.data
+            logger.info(f"RevEng.AI: Auto-unstrip status - {status}")
+
+            if status == STATUS_COMPLETED:
+                if first_poll and not resync_if_already_complete:
+                    return
+                self.call_callback(
+                    generic_return=GenericApiReturn(success=True, data=analysis_id)
+                )
+                return
+
+            if status == STATUS_FAILED:
+                self.call_callback(
+                    generic_return=GenericApiReturn(
+                        success=False,
+                        error_message="RevEng.AI auto-unstrip failed.",
+                    )
+                )
+                return
+
+            first_poll = False
+            time.sleep(POLL_INTERVAL_SECONDS)

tests/idalib/analysis_sync/test_sync.py

@@ -72,6 +72,23 @@ def test_perform_function_sync_no_matches(service, loaded_binary):
     assert result.data.total_function_count >= 1
 
 
+def test_perform_function_sync_counts_missing_symbol_names(service, auto_func):
+    func_map = FunctionMapping(
+        function_map={},
+        inverse_function_map={},
+        name_map={str(auto_func): ""},
+    )
+
+    result, pushbacks, needs_canonical = service._perform_function_sync(func_map)
+
+    assert result.success is True
+    assert result.data.missing_symbol_name_count >= 1
+    assert result.data.matched_function_count == 0
+    assert ida_name.get_name(auto_func) != ""
+    assert pushbacks == []
+    assert needs_canonical == []
+
+
 def test_perform_function_sync_classifies_invalid_name(service, auto_func):
     func_map = FunctionMapping(
         function_map={"7": auto_func},