fix: matched functions table and progress bar

Author: anisiosts

reai_toolkit/features/match_current_function/match_current_function.py

@@ -135,6 +135,7 @@ def parse_confidence(item):
                                 continue
 
                             function = bv.get_function_at(func_addr)
+                            line["function_address"] = func_addr 
                             log_info(f"RevEng.AI | Function: {function} at {func_addr:x}")
 
                             if not line["matched_function_name"] or line["matched_function_name"].startswith(("sub_", "FUN_")):
@@ -152,15 +153,14 @@ def parse_confidence(item):
                                     line["icon_text"] = "Success"
                                     matched_count += 1
 
-                            if line not in result["data"] and line["function_address"] != "N/A":
+                            if line not in result["data"]:
                                 result["data"].append(line)
 
                         except Exception as e:
                             log_error(f"RevEng.AI | Error processing function {function_by_distance.function_id}: {str(e)}")
                     #populate_table_function(result["data"])
                     #################
                     if functions_by_distance.status.lower() == "completed":
-                        
                         break
                     elif functions_by_distance.status.lower() == "error":
                         raise Exception("Function matching failed")

reai_toolkit/features/match_current_function/match_current_function_dialog.py

@@ -13,6 +13,7 @@ def __init__(self, config, match_functions, bv, func):
         super().__init__()
         self.config = config
         self.match_functions = match_functions
+        self.results = []
         self.selected_results = []
         self.bv = bv
         self.func = func
@@ -206,6 +207,8 @@ def init_ui(self):
         header.setSectionResizeMode(6, QHeaderView.Stretch)
         self.results_table.setAlternatingRowColors(True)
         self.results_table.verticalHeader().setVisible(False)
+        self.results_table.setSelectionBehavior(QTableWidget.SelectRows)
+        self.results_table.setSelectionMode(QTableWidget.SingleSelection)
 
         #functions_group_layout.addLayout(search_input_layout)
         functions_group_layout.addWidget(filter_group)
@@ -307,15 +310,16 @@ def on_matching_finished(self, success, data):
 
     def populate_results_table(self, results):
         self.selected_results.clear()
+        self.results = results
 
         self.results_table.setRowCount(0)
         self.results_table.setRowCount(len(results))
-
+        
         try:
             self.results_table.itemChanged.disconnect()
         except TypeError:
             pass  # No connections to disconnect
-
+        
         if len(results) > 0:
             self.rename_button.setEnabled(True)
             self.rename_button.setStyleSheet("""
@@ -397,6 +401,8 @@ def populate_results_table(self, results):
         self.results_table.cellClicked.connect(self.on_checkbox_changed)
 
     def on_checkbox_changed(self, item_or_row, column=None):
+        log_info(f"RevEng.AI | on_checkbox_changed: {item_or_row}, {column}")
+        
         """Handle checkbox changes to ensure only one is selected at a time"""
         if isinstance(item_or_row, QTableWidgetItem):  # Called from itemChanged
             row = item_or_row.row()
@@ -405,12 +411,14 @@ def on_checkbox_changed(self, item_or_row, column=None):
             row = item_or_row
             is_checkbox = column == 0
 
+        log_info(f"RevEng.AI | row: {row}, is_checkbox: {is_checkbox}")
+        log_info(f"RevEng.AI | len(self.results): {len(self.results)}")
         # Get the match data for this row
-        if row < len(self.current_matches):
-            match = self.current_matches[row]
+        if row < len(self.results):
+            match = self.results[row]
         else:
             return
-            
+        log_info(f"RevEng.AI | match: {match}")
         if match:
             checkbox_item = self.results_table.item(row, 0)
             current_state = checkbox_item.checkState()
@@ -432,11 +440,11 @@ def on_checkbox_changed(self, item_or_row, column=None):
                             other_checkbox.setCheckState(Qt.Unchecked)
 
                 # Update selected result with the current match
-                self.selected_result = match
-                log_info(f"RevEng.AI | Selected function match: {match.get('matched_name', 'Unknown')}")
+                self.selected_results = [match]
+                log_info(f"RevEng.AI | Selected function match: {match.get('matched_function_name', 'Unknown')}")
             else:
                 # If unchecked, clear selection
-                self.selected_result = {}
+                self.selected_result = []
                 log_info(f"RevEng.AI | Deselected function match")
 
     def start_fetching_data_types(self):
@@ -446,14 +454,14 @@ def start_fetching_data_types(self):
             self.progress = create_cancellable_progress_dialog(self, "RevEng.AI Fetch Data Types", "Fetching data types...", self.match_functions.cancel)
             self.progress.show()
             QCoreApplication.processEvents() 
-
+            """
             if not hasattr(self, 'selected_results') or not self.selected_results:
                 log_error("RevEng.AI | No current matches available for data type fetching")
                 self.progress.close()
                 QMessageBox.warning(self,"RevEng.AI Fetch Data Types","No function matches available. Please run 'Fetch Results' first.", QMessageBox.Ok)
                 return
-
-            self.fetch_data_types_thread = DataThread(self.match_functions.fetch_data_types, self.bv, self.selected_results, self.match_functions.clear_cancelled)
+            """
+            self.fetch_data_types_thread = DataThread(self.match_functions.fetch_data_types, self.bv, self.results, self.match_functions.clear_cancelled)
             self.fetch_data_types_thread.finished.connect(self.on_fetching_data_types_finished)
             self.fetch_data_types_thread.start()
 
@@ -491,6 +499,7 @@ def start_renaming(self):
         self.progress = create_progress_dialog(self, "RevEng.AI Rename Selected Functions", "Renaming Selected Functions...")
         self.progress.show()
         QCoreApplication.processEvents() 
+        log_info(f"RevEng.AI | Selected results: {len(self.selected_results)}")
 
         self.rename_thread = DataThread(
             self.match_functions.rename_functions, 

reai_toolkit/utils/core/binary_ninja.py

@@ -10,8 +10,12 @@ def rename_function(bv: BinaryView, addr: int, new_name: str, data_type: dict =
         func = bv.get_function_at(addr)
         if not func:
             log_error(f"RevEng.AI | No function found at address {hex(addr)}")
-            return False
-        
+            addr = addr + bv.image_base
+            func = bv.get_function_at(addr)
+            if not func:
+                log_error(f"RevEng.AI | No function found at address {hex(addr)}")
+                return False
+
         if func.name == new_name:
             log_info(f"RevEng.AI | Function at {hex(addr)} already has name {func.name}")
             #return False

reai_toolkit/utils/ui/binaries_popup.py

@@ -6,7 +6,7 @@
 from PySide6.QtCore import Qt, QCoreApplication
 from reai_toolkit.utils.ui.progress import create_progress_dialog
 from reai_toolkit.utils.core.threading import DataThread
-800
+
 class BinariesPopup(QDialog):
     def __init__(self, match_functions, bv, status_label=None, parent=None, write_selected_binaries=None):
         super().__init__(parent)
@@ -33,7 +33,7 @@ def _build_search_section(self, parent_layout):
         search_input_layout = QHBoxLayout()
         self.search_input = QLineEdit()
         self.search_input.setPlaceholderText("Enter search term...")
-        self.search_input.returnPressed.connect(self._search_binaries)
+        #self.search_input.returnPressed.connect(self._search_binaries)
 
         description_label = QLabel(
             "Search (e.g. sha_256_hash:{}, tag:{}, binary_name:{}, function_name:{}, model_name:{})"

reai_toolkit/utils/ui/collections_popup.py

@@ -6,7 +6,7 @@
 from PySide6.QtCore import Qt, QCoreApplication
 from reai_toolkit.utils.ui.progress import create_progress_dialog
 from reai_toolkit.utils.core.threading import DataThread
-800
+
 class CollectionsPopup(QDialog):
     def __init__(self, match_functions, bv, status_label=None, parent=None, write_selected_collections=None):
         super().__init__(parent)
@@ -33,7 +33,7 @@ def _build_search_section(self, parent_layout):
         search_input_layout = QHBoxLayout()
         self.search_input = QLineEdit()
         self.search_input.setPlaceholderText("Enter search term...")
-        self.search_input.returnPressed.connect(self._search_collections)
+        #self.search_input.returnPressed.connect(self._search_collections)
 
         description_label = QLabel(
             "Search (e.g. sha_256_hash:{}, tag:{}, collection_name:{}, function_name:{}, model_name:{})"