fix: analysis sync

Author: vgoat21

reai_toolkit/features/choose_source/choose_source.py

@@ -1,11 +1,12 @@
 import revengai
 from reai_toolkit.utils import get_sha256
 from reai_toolkit.utils.core.sync import AnalysisSyncService
+from reai_toolkit.features.configuration.config import Config
 from binaryninja import BinaryView, log_info, log_error
 
 class ChooseSource:
-    def __init__(self, config):
-        self.config = config
+    def __init__(self, config: Config):
+        self.config: Config = config
 
         self.sync_service = AnalysisSyncService(config)
 

reai_toolkit/features/configuration/config.py

@@ -3,6 +3,7 @@
 from reai_toolkit.utils import get_sha256
 from binaryninja.interaction import InteractionHandler
 from binaryninja import Settings, log_info, log_error, BinaryView
+from reai_toolkit.utils.core.sync import AnalysisSyncService
 
 
 class Config:
@@ -145,6 +146,7 @@ def init_config(self, bv: BinaryView):
                 self.binary_id = all_analyses[self.sha256]["binary_id"]
                 self.analysis_id = all_analyses[self.sha256]["analysis_id"]
                 self.model_id = all_analyses[self.sha256]["model_id"]
+                AnalysisSyncService(self).sync_analysis_data(analysis_id=self.analysis_id, bv=bv)
             else:
                 log_info(f"RevEng.AI | Binary not found in saved configurations, searching in RevEng.AI...")
                 with revengai.ApiClient(self.api_config) as api_client:
@@ -156,6 +158,7 @@ def init_config(self, bv: BinaryView):
                     self.binary_id = api_response.data.results[0].binary_id
                     self.analysis_id = api_response.data.results[0].analysis_id
                     self.model_id = api_response.data.results[0].model_id
+                    AnalysisSyncService(self).sync_analysis_data(analysis_id=self.analysis_id, bv=bv)
                     log_info(f"RevEng.AI | Binary found in RevEng.AI, binary_id: {self.binary_id}")
                     self.set_current_info(self.binary_id, self.analysis_id, self.model_id)
 

reai_toolkit/features/upload/upload.py

@@ -95,7 +95,7 @@ def upload_binary(self, bv: BinaryView, options: dict):
 
                 log_info(f"RevEng.AI | Analysis started successfully. Analysis ID: {analysis_result.data.analysis_id}, Binary ID: {analysis_result.data.binary_id}")
 
-            PeriodicChecker().start_checking(bv, analysis_result.data.analysis_id, analysis_result.data.binary_id, self.config.set_current_info, self.config.api_config)
+            PeriodicChecker(self.config).start_checking(bv, analysis_result.data.analysis_id, analysis_result.data.binary_id, self.config.set_current_info, self.config.api_config)
 
             return True, "Analysis started successfully."
 

reai_toolkit/utils/core/sync.py

@@ -7,8 +7,9 @@ class AnalysisSyncService:
 
     sdk_config: Configuration
 
-    def __init__(self, sdk_config: Configuration):
-        self.sdk_config = sdk_config
+    def __init__(self, config):
+        self.config = config
+        self.sdk_config = config.api_config
 
     def _get_current_base_address(self, bv) -> int:
         return bv.start
@@ -46,7 +47,6 @@ def _fetch_function_map(self, analysis_id: int) -> FunctionMapping:
                 analysis_id=analysis_id
             )
             func_map = function_map.data.function_maps
-            self.safe_put_function_mapping(func_map=func_map)
             return func_map
 
     def _match_functions(
@@ -108,15 +108,9 @@ def sync_analysis_data(
         """
         Syncs the analysis data until completion or failure.
         """
-        response = self.api_request_returning(
-            fn=lambda: self._fetch_function_map(analysis_id=analysis_id)
-        )
-
-        if not response.success:
-            self.call_callback(generic_return=response)
-            return
+        response = self._fetch_function_map(analysis_id=analysis_id)
 
-        function_mapping: FunctionMapping = response.data
+        function_mapping: FunctionMapping = response
 
         self._match_functions(func_map=function_mapping, bv=bv)
 

reai_toolkit/utils/monitoring/process_binary_monitor.py

@@ -6,15 +6,18 @@
 from requests.exceptions import RequestException
 from PySide6.QtWidgets import QMessageBox
 from PySide6.QtCore import QObject, Signal
+from reai_toolkit.utils.core.sync import AnalysisSyncService
 
 class PeriodicChecker(QObject):
     update_text_signal = Signal(object, str)
+    sync_service: AnalysisSyncService
 
-    def __init__(self):
+    def __init__(self, config):
         super().__init__()
         self._current_timer: Optional[Timer] = None
         self.number_of_clicks = 0
         self.update_text_signal.connect(self._update_text_slot)
+        self.sync_service = AnalysisSyncService(config)
 
     def _update_text_slot(self, callback, text):
         """Slot that runs in the main thread to safely update UI"""
@@ -62,6 +65,9 @@ def _worker(bv: BinaryView, bid: int, aid: int):
                         )
                         model_id = analysis_details.data.model_id
                         callback(bid, aid, model_id)
+
+                        self.sync_service.sync_analysis_data(analysis_id=aid, bv=bv)
+
                         log_info(f"RevEng.AI | Analysis completed with status: {status} for Binary ID: {bid} | Analysis ID: {aid} | Model ID: {model_id}")
             except RequestException as ex:
                 log_error(f"RevEng.AI | Error getting binary analysis status: {str(ex)}")