From ac70f2a253832ab3ff03822c75f05db608aed7b0 Mon Sep 17 00:00:00 2001 From: Joshua Villalta <93329405+vgoat21@users.noreply.github.com> Date: Fri, 10 Jul 2026 13:28:24 +0100 Subject: [PATCH] feat(PLU-301): auto unstrip re-sync --- reai_toolkit/app/app.py | 6 ++ reai_toolkit/app/coordinator.py | 14 ++++ .../coordinators/auto_unstrip_coordinator.py | 69 +++++++++++++++ .../app/coordinators/detach_coordinator.py | 1 + .../existing_analysis_coordinator.py | 10 +++ .../coordinators/poll_status_coordinator.py | 12 +++ .../coordinators/sync_analysis_coordinator.py | 5 ++ .../services/analysis_sync/analysis_sync.py | 5 ++ .../app/services/analysis_sync/schema.py | 1 + .../services/auto_unstrip_status/__init__.py | 0 .../auto_unstrip_status.py | 84 +++++++++++++++++++ tests/idalib/analysis_sync/test_sync.py | 17 ++++ 12 files changed, 224 insertions(+) create mode 100644 reai_toolkit/app/coordinators/auto_unstrip_coordinator.py create mode 100644 reai_toolkit/app/services/auto_unstrip_status/__init__.py create mode 100644 reai_toolkit/app/services/auto_unstrip_status/auto_unstrip_status.py diff --git a/reai_toolkit/app/app.py b/reai_toolkit/app/app.py index a2088ed2..e1d2733b 100644 --- a/reai_toolkit/app/app.py +++ b/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 ) diff --git a/reai_toolkit/app/coordinator.py b/reai_toolkit/app/coordinator.py index b92ce29a..7527b70b 100644 --- a/reai_toolkit/app/coordinator.py +++ b/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, ) ) diff --git a/reai_toolkit/app/coordinators/auto_unstrip_coordinator.py b/reai_toolkit/app/coordinators/auto_unstrip_coordinator.py new file mode 100644 index 00000000..6cd1ff48 --- /dev/null +++ b/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 + ) diff --git a/reai_toolkit/app/coordinators/detach_coordinator.py b/reai_toolkit/app/coordinators/detach_coordinator.py index 3b8cde2c..a56f8e2c 100644 --- a/reai_toolkit/app/coordinators/detach_coordinator.py +++ b/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() diff --git a/reai_toolkit/app/coordinators/existing_analysis_coordinator.py b/reai_toolkit/app/coordinators/existing_analysis_coordinator.py index 3b00a225..ef9439d1 100644 --- a/reai_toolkit/app/coordinators/existing_analysis_coordinator.py +++ b/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() diff --git a/reai_toolkit/app/coordinators/poll_status_coordinator.py b/reai_toolkit/app/coordinators/poll_status_coordinator.py index e8338146..0d15f415 100644 --- a/reai_toolkit/app/coordinators/poll_status_coordinator.py +++ b/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() diff --git a/reai_toolkit/app/coordinators/sync_analysis_coordinator.py b/reai_toolkit/app/coordinators/sync_analysis_coordinator.py index 0c08650e..17b05d05 100644 --- a/reai_toolkit/app/coordinators/sync_analysis_coordinator.py +++ b/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 " diff --git a/reai_toolkit/app/services/analysis_sync/analysis_sync.py b/reai_toolkit/app/services/analysis_sync/analysis_sync.py index 92248d8a..9afa1219 100644 --- a/reai_toolkit/app/services/analysis_sync/analysis_sync.py +++ b/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,6 +178,8 @@ 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 @@ -184,10 +187,12 @@ def _perform_function_sync( 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, ) diff --git a/reai_toolkit/app/services/analysis_sync/schema.py b/reai_toolkit/app/services/analysis_sync/schema.py index 96094b66..d6c7c16f 100644 --- a/reai_toolkit/app/services/analysis_sync/schema.py +++ b/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 diff --git a/reai_toolkit/app/services/auto_unstrip_status/__init__.py b/reai_toolkit/app/services/auto_unstrip_status/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/reai_toolkit/app/services/auto_unstrip_status/auto_unstrip_status.py b/reai_toolkit/app/services/auto_unstrip_status/auto_unstrip_status.py new file mode 100644 index 00000000..a988e1b2 --- /dev/null +++ b/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) diff --git a/tests/idalib/analysis_sync/test_sync.py b/tests/idalib/analysis_sync/test_sync.py index 0e6f07f6..427a446e 100644 --- a/tests/idalib/analysis_sync/test_sync.py +++ b/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},