feat(PLU-210): sync functions - saved

Author: vgoat21

reai_toolkit/features/choose_source/choose_source.py

@@ -1,11 +1,14 @@
 import revengai
 from reai_toolkit.utils import get_sha256
+from reai_toolkit.utils.core.sync import AnalysisSyncService
 from binaryninja import BinaryView, log_info, log_error
 
 class ChooseSource:
     def __init__(self, config):
         self.config = config
 
+        self.sync_service = AnalysisSyncService(config)
+
     def choose_source(self, bv: BinaryView, chose: str): 
         try:
             log_info(f"RevEng.AI | Item: {chose}")
@@ -23,6 +26,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/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,7 +88,7 @@ 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
                 )

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, sdk_config: Configuration):
+        self.sdk_config = sdk_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
+            self.safe_put_function_mapping(func_map=func_map)
+            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
+
+                # 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.api_request_returning(
+            fn=lambda: self._fetch_function_map(analysis_id=analysis_id)
+        )
+
+        if not response.success:
+            self.call_callback(generic_return=response)
+            return
+
+        function_mapping: FunctionMapping = response.data
+
+        self._match_functions(func_map=function_mapping, bv=bv)
+
+        self._fetch_basic_and_rebase(bv=bv, analysis_id=analysis_id)