fix: function matching (batch) api update

Author: vgoat21

reai_toolkit/features/match_functions/match_functions.py

@@ -137,7 +137,9 @@ def parse_confidence(item):
                                 "icon_text": "Failed",
                                 "function_address": function_by_distance.matched_functions[0].function_vaddr,
                                 "function_name": "N/A",
+                                "source_function_id": function_by_distance.function_id,
                                 "matched_function_name": function_by_distance.matched_functions[0].function_name,
+                                "matched_mangled_name": function_by_distance.matched_functions[0].mangled_name,
                                 "signature": "N/A",
                                 "matched_hash": function_by_distance.matched_functions[0].sha_256_hash,
                                 "matched_binary_name": function_by_distance.matched_functions[0].binary_name,
@@ -219,7 +221,7 @@ def parse_confidence(item):
             raise e
 
 
-    def _process_rename_batch(self, chunk: List[Dict], bv: BinaryView, deci: DecompilerInterface = None) -> Tuple[int, int]:
+    def _process_rename_batch(self, config, chunk: List[Dict], bv: BinaryView, deci: DecompilerInterface = None) -> Tuple[int, int]:
         try:
             log_info(f"RevEng.AI | Processing chunk of {len(chunk)} functions")
             renamed_count = 0
@@ -230,7 +232,7 @@ def _process_rename_batch(self, chunk: List[Dict], bv: BinaryView, deci: Decompi
                         return 0, 0
 
                     addr = int(result['function_address'])
-                    if rename_function_util(bv, addr, result["matched_function_name"]):
+                    if rename_function_util(config, bv, addr, result["matched_function_name"], result["matched_mangled_name"], result["source_function_id"]):
                         renamed_count += 1
 
                         if result.get('signature_data', None) is not None:
@@ -270,7 +272,7 @@ def rename_functions(self, bv: BinaryView, selected_results: List[Dict]) -> List
 
             with ThreadPoolExecutor(max_workers=4) as executor:
                 future_to_chunk = {
-                    executor.submit(self._process_rename_batch, chunk, bv, deci): i 
+                    executor.submit(self._process_rename_batch, self.config, chunk, bv, deci): i 
                     for i, chunk in enumerate(chunks)
                 }
 

reai_toolkit/features/match_functions/match_functions_dialog.py

@@ -189,21 +189,22 @@ def init_ui(self):
 
 
         self.results_table = QTableWidget()
-        self.results_table.setColumnCount(8)
+        self.results_table.setColumnCount(9)
         self.results_table.setHorizontalHeaderLabels([
-            "Virtual Address", "Function Name", "Matched Function", "Signature", "Similarity", "Confidence", "Matched Hash", "Matched Binary"
+            "Select", "Virtual Address", "Function Name", "Matched Function", "Signature", "Similarity", "Confidence", "Matched Hash", "Matched Binary"
         ])
         self.results_table.setSelectionMode(QAbstractItemView.NoSelection)
 
         header = self.results_table.horizontalHeader()
         header.setSectionResizeMode(0, QHeaderView.ResizeToContents)
-        header.setSectionResizeMode(1, QHeaderView.Stretch)
+        header.setSectionResizeMode(1, QHeaderView.ResizeToContents)
         header.setSectionResizeMode(2, QHeaderView.Stretch)
-        header.setSectionResizeMode(3, QHeaderView.ResizeToContents)
+        header.setSectionResizeMode(3, QHeaderView.Stretch)
         header.setSectionResizeMode(4, QHeaderView.ResizeToContents)
         header.setSectionResizeMode(5, QHeaderView.ResizeToContents)
-        header.setSectionResizeMode(6, QHeaderView.Stretch)
+        header.setSectionResizeMode(6, QHeaderView.ResizeToContents)
         header.setSectionResizeMode(7, QHeaderView.Stretch)
+        header.setSectionResizeMode(8, QHeaderView.Stretch)
         self.results_table.setAlternatingRowColors(True)
         self.results_table.verticalHeader().setVisible(False)
 
@@ -306,7 +307,8 @@ def on_matching_finished(self, success, data):
 
 
     def populate_results_table(self, results):
-        self.selected_results.clear()
+        # Store results first so _update_selected_results can access them
+        self.all_results = results
 
         self.results_table.setRowCount(0)
         self.results_table.setRowCount(len(results))
@@ -336,6 +338,21 @@ def populate_results_table(self, results):
             """)
 
         for row, match in enumerate(results):
+            # Add checkbox in first column
+            checkbox = QCheckBox()
+            if match.get("icon_text", "Failed") == "Success":
+                checkbox.setChecked(True)
+            else:
+                checkbox.setChecked(False)
+            checkbox.stateChanged.connect(lambda state, m=match: self._on_checkbox_changed(state, m))
+            
+            checkbox_widget = QWidget()
+            checkbox_layout = QHBoxLayout(checkbox_widget)
+            checkbox_layout.addWidget(checkbox)
+            checkbox_layout.setAlignment(Qt.AlignCenter)
+            checkbox_layout.setContentsMargins(0, 0, 0, 0)
+            self.results_table.setCellWidget(row, 0, checkbox_widget)
+            
             column_data = [
                 "function_address",
                 "function_name",
@@ -347,7 +364,7 @@ def populate_results_table(self, results):
                 "matched_binary_name"
             ] 
 
-            for column, field in enumerate(column_data, start=0):
+            for column, field in enumerate(column_data, start=1):
                 value = match.get(field, "N/A")
                 if field == "function_address":
                     if isinstance(value, int):
@@ -373,17 +390,38 @@ def populate_results_table(self, results):
                         icon_path = f"{os.path.dirname(__file__)}/../../images/success.png"
                         item.setToolTip(value)
                     item.setIcon(QIcon(icon_path))
-                    #item.setIconAlignment(Qt.AlignCenter)
                     item.setText("")
                 else:
                     item.setToolTip(value)
                 self.results_table.setItem(row, column, item)
-            
-            if match.get("icon_text", "Failed") == "Success":
-                self.selected_results.append(match)
+        
+        # Initialize selected_results based on checkboxes
+        self._update_selected_results()
+    
+    def _on_checkbox_changed(self, state, match):
+        self._update_selected_results()
+    
+    def _update_selected_results(self):
+        self.selected_results.clear()
+        for row in range(self.results_table.rowCount()):
+            checkbox_widget = self.results_table.cellWidget(row, 0)
+            if checkbox_widget:
+                checkbox = checkbox_widget.findChild(QCheckBox)
+                if checkbox and checkbox.isChecked():
+                    # Get the match data from the row
+                    addr_item = self.results_table.item(row, 1)
+                    if addr_item:
+                        addr_text = addr_item.text()
+                        # Find the corresponding match in all_results
+                        if hasattr(self, 'all_results'):
+                            for match in self.all_results:
+                                match_addr = match.get("function_address", "")
+                                if isinstance(match_addr, int):
+                                    match_addr = f"0x{match_addr:x}"
+                                if match_addr == addr_text:
+                                    self.selected_results.append(match)
+                                    break
         log_info(f"RevEng.AI | Selected results: {len(self.selected_results)}")
-        for result in self.selected_results:
-            log_info(f"RevEng.AI | Result: {result}")
 
 
     def start_fetching_data_types(self):

reai_toolkit/utils/core/binary_ninja.py

@@ -5,7 +5,24 @@
 from os.path import isfile
 from os import access, R_OK
 
-def rename_function(bv: BinaryView, addr: int, new_name: str, data_type: dict = None) -> bool:
+
+def _rename_in_portal(config: revengai.Configuration, function_id:int, new_name:str, new_mangled_name:str):
+    try:
+        with revengai.ApiClient(config.api_config) as api_client:
+            api_instance = revengai.FunctionsRenamingHistoryApi(api_client)
+            api_instance.rename_function_id(
+                function_id=function_id,
+                function_rename=revengai.FunctionRename(
+                    new_name=new_name,
+                    new_mangled_name=new_mangled_name
+                )
+            )
+            log_info(f"RevEng.AI | Renamed function in portal at {function_id} to {new_name}")
+
+    except Exception as e:
+        log_error(f"RevEng.AI | Error renaming function in portal at {function_id}: {str(e)}")
+
+def rename_function(config, bv: BinaryView, addr: int, new_name: str, new_mangled_name: str, source_function_id: int, data_type: dict = None) -> bool:
     try:
         func = bv.get_function_at(addr)
         if not func:
@@ -22,6 +39,8 @@ def rename_function(bv: BinaryView, addr: int, new_name: str, data_type: dict =
 
         new_symbol = Symbol(SymbolType.FunctionSymbol, addr, new_name)
         bv.define_user_symbol(new_symbol)
+
+        _rename_in_portal(config, source_function_id, new_name, new_mangled_name)
 
         log_info(f"RevEng.AI | Renamed function at {hex(addr)} to {new_name}")
         return True