Merge pull request #8 from RevEngAI/feat-PLU-210

feat(PLU-210): sync functions

Author: vgoat21

reai_toolkit/features/choose_source/choose_source.py

@@ -1,10 +1,14 @@
 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)
 
     def choose_source(self, bv: BinaryView, chose: str): 
         try:
@@ -23,6 +27,8 @@ def choose_source(self, bv: BinaryView, chose: str):
             log_info(f"RevEng.AI | Changing Model ID to {new_model_id}")
             self.config.set_current_info(new_binary_id, new_analysis_id, new_model_id)
 
+            self.sync_service.sync_analysis_data(analysis_id=new_analysis_id, bv=bv)
+
             return True, "Binary ID changed successfully."
         except Exception as e:
             log_error(f"RevEng.AI | Failed to choose source: {str(e)}")

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

@@ -3,6 +3,7 @@
 from binaryninja import BinaryView, log_info, log_error
 from os.path import basename
 from reai_toolkit.utils import PeriodicChecker
+from reai_toolkit.utils.core.sync import AnalysisSyncService
 
 class BinaryUploader:
     def __init__(self, config):
@@ -87,14 +88,14 @@ def upload_binary(self, bv: BinaryView, options: dict):
                     analysis_scope=revengai.AnalysisScope.PRIVATE if options["is_private"] else revengai.AnalysisScope.PUBLIC,
                     symbols=symbols
                 )
-                
+
                 analysis_result = analyses_client.create_analysis(
                     analysis_create_request=analysis_create_request
                 )
 
                 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/binary_ninja.py

@@ -37,7 +37,7 @@ def rename_function(config, bv: BinaryView, addr: int, new_name: str, new_mangle
             log_info(f"RevEng.AI | Function at {hex(addr)} already has name {func.name}")
             #return False
 
-        new_symbol = Symbol(SymbolType.FunctionSymbol, addr, new_name)
+        new_symbol = Symbol(SymbolType.FunctionSymbol, addr, new_mangled_name)
         bv.define_user_symbol(new_symbol)
 
         _rename_in_portal(config, source_function_id, new_name, new_mangled_name)

reai_toolkit/utils/core/sync.py

@@ -0,0 +1,123 @@
+from binaryninja import log_info, BinaryView
+from revengai import AnalysesCoreApi, Configuration, FunctionMapping
+from revengai import BaseResponseBasic, AnalysesCoreApi, ApiClient
+from binaryninja import BinaryView, log_error, log_info, Symbol, SymbolType, Function
+
+class AnalysisSyncService:
+
+    sdk_config: Configuration
+
+    def __init__(self, config):
+        self.config = config
+        self.sdk_config = config.api_config
+
+    def _get_current_base_address(self, bv) -> int:
+        return bv.start
+
+    def _rebase_program(self, bv, base_address_delta: int) -> None:
+        bv.rebase(bv.start + base_address_delta)
+
+    def _fetch_basic_and_rebase(self, bv: BinaryView, analysis_id: int) -> BaseResponseBasic:
+        """
+        Fetches basic analysis information and rebases the program if necessary.
+        """
+        with ApiClient(self.sdk_config) as api_client:
+            analyses_client = AnalysesCoreApi(api_client)
+            analysis_details: BaseResponseBasic = analyses_client.get_analysis_basic_info(
+                analysis_id=analysis_id
+            )
+
+            local_base_address: int = self._get_current_base_address(bv)
+
+            if analysis_details.data and analysis_details.data.base_address is not None:
+                remote_base_address: int = analysis_details.data.base_address
+
+            if local_base_address != remote_base_address:
+                    base_address_delta: int = remote_base_address - local_base_address
+                    self._rebase_program(bv, base_address_delta)
+
+    def _fetch_function_map(self, analysis_id: int) -> FunctionMapping:
+        """
+        Fetches the function map for the given analysis ID.
+        """
+        with ApiClient(self.sdk_config) as api_client:
+            analyses_client = AnalysesCoreApi(api_client)
+
+            function_map = analyses_client.get_analysis_function_map(
+                analysis_id=analysis_id
+            )
+            func_map = function_map.data.function_maps
+            return func_map
+
+    def _match_functions(
+        self,
+        func_map: FunctionMapping,
+        bv: BinaryView,
+    ) -> None:
+        function_map = func_map.function_map
+        inverse_function_map = func_map.inverse_function_map
+
+        log_info(
+            f"RevEng.AI | Retrieved {len(function_map)} function mappings from analysis"
+        )
+
+        # Compute which IDA functions match the revengai analysis functions
+        matched_functions = []
+        unmatched_local_functions = []
+        unmatched_remote_functions = []
+
+        # Track local functions matched
+        local_function_vaddrs_matched = set()
+
+        for func in bv.functions:
+            start_ea = func.start
+            if str(start_ea) in inverse_function_map:
+                new_name: str | None = func_map.name_map.get(str(start_ea), None)
+                if new_name is None:
+                    continue
+
+                # Check if function has a user-defined symbol, skip if it does
+                if func.symbol and func.symbol.auto == False:
+                    log_info(f"RevEng.AI | Skipping user-defined function at 0x{start_ea:x}: {func.name}")
+                    local_function_vaddrs_matched.add(start_ea)
+                    continue
+
+                # Rename local function
+                new_symbol = Symbol(SymbolType.FunctionSymbol, start_ea, new_name)
+                bv.define_user_symbol(new_symbol)
+                
+                matched_functions.append(
+                    (int(inverse_function_map[str(start_ea)]), start_ea)
+                )
+                local_function_vaddrs_matched.add(start_ea)
+            else:
+                unmatched_local_functions.append(start_ea)
+
+        unmatched_portal_map = {}
+        # Track remote functions not matched
+        for func_id_str, func_vaddr in function_map.items():
+            if int(func_vaddr) not in local_function_vaddrs_matched:
+                unmatched_remote_functions.append((int(func_vaddr), int(func_id_str)))
+                unmatched_portal_map[int(func_vaddr)] = int(func_id_str)
+
+        log_info(f"RevEng.AI | Matched {len(matched_functions)} functions")
+        log_info(
+            f"RevEng.AI | {len(unmatched_local_functions)} local functions not matched"
+        )
+        log_info(
+            f"RevEng.AI | {len(unmatched_remote_functions)} remote functions not matched"
+        )
+
+    def sync_analysis_data(
+        self, analysis_id: int, bv: BinaryView
+    ) -> None:
+        """
+        Syncs the analysis data until completion or failure.
+        """
+        response = self._fetch_function_map(analysis_id=analysis_id)
+
+        function_mapping: FunctionMapping = response
+
+        self._match_functions(func_map=function_mapping, bv=bv)
+
+        self._fetch_basic_and_rebase(bv=bv, analysis_id=analysis_id)

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)}")