improving search tab

Author: anisiosts

features/match_functions/match_functions.py

@@ -11,14 +11,16 @@ def __init__(self, config):
         self.binary_id = None
         self.analyzed_functions = []
 
-    def search_collections(self, search_term: str = "") -> List[Dict]:
-        """Search for collections in RevEng.AI database"""
+    def search_collections(self, bv: BinaryView, search_term: str = "") -> List[Dict]:
         try:
             log_info(f"RevEng.AI | Searching collections with term: '{search_term}'")
+            query = self._parse_search_query(search_term)
+            log_info(f"RevEng.AI | Query: {query}")
 
             # Since RE_collections might not be available, we'll use RE_search to find binaries
             # and simulate collections based on search results
-            search_results = RE_search(fpath=search_term if search_term else "").json()
+            """
+            search_results = RE_search(fpath=bv.file.filename, search_term=search_term).json()
             
             if "query_results" not in search_results:
                 log_error("RevEng.AI | No search results found")
@@ -42,10 +44,10 @@ def search_collections(self, search_term: str = "") -> List[Dict]:
             
             log_info(f"RevEng.AI | Found {len(collections)} collections")
             return collections
-            
+            """            
         except Exception as e:
             log_error(f"RevEng.AI | Error searching collections: {str(e)}")
-            return []
+            return False, str(e)
 
     def get_collection_functions(self, collection_id: str) -> List[Dict]:
         """Get functions from a specific collection (simulated)"""
@@ -180,4 +182,103 @@ def get_function_details(self, bv: BinaryView, function_address: int) -> Optiona
 
         except Exception as e:
             log_error(f"RevEng.AI | Error getting function details: {str(e)}")
-            return None 
+            return None 
+        
+    def _parse_search_query(self, query):
+        """
+        Parse a search query with special selectors.
+
+        Args:
+            query (str): The search query string to parse
+
+        Returns:
+            dict: A dictionary containing parsed query components
+
+        Raises:
+            ValueError: If multiple non-tag selectors or a selector with raw
+                        query are used
+        """
+        # Initialize the result dictionary with default empty values
+        result = {
+            'query': None,
+            'sha_256_hash': None,
+            'tags': None,
+            'binary_name': None,
+            'collection_name': None,
+            'function_name': None,
+            'model_name': None
+        }
+
+        # List of possible selectors (excluding 'tag')
+        single_selectors = [
+            'sha_256_hash',
+            'binary_name',
+            'collection_name',
+            'function_name',
+            'model_name'
+        ]
+
+        # Parse selector-based queries
+        def extract_selector_value(query, selector):
+            """Helper function to extract selector value"""
+            selector_pattern = f"{selector}:"
+            selector_match = query.find(selector_pattern)
+
+            if selector_match != -1:
+                # Extract the value after the selector
+                start = selector_match + len(selector_pattern)
+                end = query.find(' ', start)
+
+                # If no space found, take till the end of string
+                if end == -1:
+                    end = len(query)
+
+                # Extract the value and the full selector part
+                value = query[start:end].strip()
+                full_selector_part = query[selector_match:end].strip()
+
+                return value, full_selector_part
+
+            return None, None
+
+        # Process tags first (can be multiple)
+        def process_tags(query):
+            tags = []
+            while True:
+                tag_value, tag_part = extract_selector_value(query, 'tag')
+                if not tag_value:
+                    break
+                tags.append(tag_value)
+                query = query.replace(tag_part, '').strip()
+            if len(tags) == 0:
+                tags = None
+            return tags, query
+
+        # Process tags
+        result['tags'], query = process_tags(query)
+
+        # Process other single selectors
+        for selector in single_selectors:
+            value, selector_part = extract_selector_value(query, selector)
+
+            if value:
+                # Check if this selector was already set
+                if result[selector] is not None:
+                    raise ValueError(
+                        f"Only one {selector} selector can be used.")
+
+                result[selector] = value
+                query = query.replace(selector_part, '').strip()
+
+        # Validation checks for additional text
+        query = query.strip()
+        if query:
+            # If query is not empty after removing selectors
+            if any(result[selector] is not None for selector in
+                   single_selectors):
+                raise ValueError(
+                    "Selector cannot be used with additional text.")
+            # If no other selectors, treat as raw query
+            result['query'] = query
+
+        return result

features/match_functions/match_functions_dialog.py

@@ -3,11 +3,12 @@
                              QLabel, QLineEdit, QTableWidget, QTableWidgetItem,
                              QHeaderView, QTabWidget, QWidget, QMessageBox,
                              QCheckBox, QDoubleSpinBox, QSpinBox, QGroupBox,
-                             QSplitter, QTextEdit, QProgressBar)
-from PySide6.QtCore import Qt, QTimer
+                             QSplitter, QTextEdit, QProgressBar, QSlider)
+from PySide6.QtCore import Qt, QTimer, QCoreApplication
 from PySide6.QtGui import QIcon
 from revengai_bn.utils import create_progress_dialog
 from .match_functions_thread import MatchFunctionsThread
+from .search_collections_thread import SearchCollectionsThread
 import os
 
 class MatchFunctionsDialog(QDialog):
@@ -53,14 +54,21 @@ def create_search_tab(self):
         layout = QVBoxLayout()
 
         # Search section
-        search_group = QGroupBox("Search Collections")
+        #search_group = QGroupBox("Search Collections")
+        search_group = QGroupBox()
         search_layout = QVBoxLayout()
 
         # Search input
         search_input_layout = QHBoxLayout()
         self.search_input = QLineEdit()
-        self.search_input.setPlaceholderText("Enter search term (e.g., 'stealc')")
+        self.search_input.setPlaceholderText("Enter search term")
         self.search_input.returnPressed.connect(self.search_collections)
+        description_label = QLabel(
+            "Search (e.g. sha_256_hash:{}, tag:{}, binary_name:{}, collection_name:{}, function_name:{}, model_name:{})"
+            #"Search Syntax: sha_256_hash: {hash}, tag: {tag}, binary_name: {binary_name}, collection_name: {collection_name},function_name: {function_name}, model_name: {model_name}"
+        )
+        description_label.setWordWrap(True)
+        
 
         self.search_button = QPushButton("Search")
         self.search_button.clicked.connect(self.search_collections)
@@ -70,6 +78,7 @@ def create_search_tab(self):
                 color: white;
                 padding: 6px 12px;
                 border-radius: 4px;
+                
             }
             QPushButton:hover {
                 background-color: #0056b3;
@@ -78,7 +87,9 @@ def create_search_tab(self):
 
         search_input_layout.addWidget(self.search_input)
         search_input_layout.addWidget(self.search_button)
+       
         search_layout.addLayout(search_input_layout)
+        search_layout.addWidget(description_label)
 
         # Collections table
         self.collections_table = QTableWidget()
@@ -94,36 +105,38 @@ def create_search_tab(self):
 
         self.collections_table.setSelectionBehavior(QTableWidget.SelectRows)
         self.collections_table.setAlternatingRowColors(True)
+        self.collections_table.itemSelectionChanged.connect(self.on_result_selection_changed)
 
         search_layout.addWidget(self.collections_table)
         search_group.setLayout(search_layout)
         layout.addWidget(search_group)
 
         # Match settings section
-        settings_group = QGroupBox("Match Settings")
+        #settings_group = QGroupBox("Match Settings")
+        settings_group = QGroupBox()
         settings_layout = QVBoxLayout()
-        
-        # Distance threshold
-        distance_layout = QHBoxLayout()
-        distance_layout.addWidget(QLabel("Distance Threshold:"))
-        self.distance_spinbox = QDoubleSpinBox()
-        self.distance_spinbox.setRange(0.01, 1.0)
-        self.distance_spinbox.setValue(0.1)
-        self.distance_spinbox.setSingleStep(0.01)
-        self.distance_spinbox.setDecimals(2)
-        distance_layout.addWidget(self.distance_spinbox)
-        distance_layout.addStretch()
-        settings_layout.addLayout(distance_layout)
-        
-        # Max matches
-        matches_layout = QHBoxLayout()
-        matches_layout.addWidget(QLabel("Max Matches:"))
-        self.max_matches_spinbox = QSpinBox()
-        self.max_matches_spinbox.setRange(1, 100)
-        self.max_matches_spinbox.setValue(10)
-        matches_layout.addWidget(self.max_matches_spinbox)
-        matches_layout.addStretch()
-        settings_layout.addLayout(matches_layout)
+
+        # Confidence slider
+        confidence_layout = QHBoxLayout()
+        confidence_layout.addWidget(QLabel("Confidence:"))
+        self.confidenceSlider = QSlider()
+        self.confidenceSlider.setMaximum(100)
+        self.confidenceSlider.setPageStep(5)
+        self.confidenceSlider.setSliderPosition(90)
+        self.confidenceSlider.setOrientation(Qt.Horizontal)
+        self.confidenceSlider.setInvertedAppearance(False)
+        self.confidenceSlider.setInvertedControls(False)
+        self.confidenceSlider.setTickPosition(QSlider.TicksBothSides)
+        self.confidenceSlider.setTickInterval(5)
+        self.confidenceSlider.setObjectName("confidenceSlider")
+        confidence_layout.addWidget(self.confidenceSlider)
+        
+        # Add confidence value label
+        self.confidence_value_label = QLabel("90")
+        self.confidenceSlider.valueChanged.connect(lambda value: self.confidence_value_label.setText(str(value)))
+        confidence_layout.addWidget(self.confidence_value_label)
+        
+        settings_layout.addLayout(confidence_layout)
 
         # Debug symbols checkbox
         self.debug_symbols_checkbox = QCheckBox("Limit Matches to Debug Symbols")
@@ -136,27 +149,33 @@ def create_search_tab(self):
         # Action buttons
         button_layout = QHBoxLayout()
 
-        self.match_button = QPushButton("Match Functions")
-        self.match_button.setStyleSheet("""
-            QPushButton {
-                background-color: #28a745;
-                color: white;
-                padding: 8px 16px;
-                border-radius: 4px;
-                font-weight: bold;
-            }
-            QPushButton:hover {
-                background-color: #218838;
-            }
-            QPushButton:disabled {
-                background-color: #6c757d;
-            }
-        """)
-        self.match_button.clicked.connect(self.start_matching)
+        self.fetch_results_button = QPushButton("Fetch Results")
+        self.fetch_data_types_button = QPushButton("Fetch Data Types") 
+        self.rename_selected_button = QPushButton("Rename Selected")
+
+        for button in [self.fetch_results_button, self.fetch_data_types_button, self.rename_selected_button]:
+            button.setStyleSheet("""
+                QPushButton {
+                    background-color: #6c757d;
+                    color: white;
+                    padding: 6px 12px;
+                    border-radius: 4px;
+                }
+                QPushButton:hover {
+                    background-color: #5a6268;
+                }
+            """)
+
+        self.fetch_results_button.clicked.connect(self.start_matching)
+        self.fetch_data_types_button.clicked.connect(self.start_matching)
+        self.rename_selected_button.clicked.connect(self.start_matching)
 
-        button_layout.addWidget(self.match_button)
+        button_layout.addWidget(self.fetch_results_button)
+        button_layout.addWidget(self.fetch_data_types_button)
+        button_layout.addWidget(self.rename_selected_button)
         button_layout.addStretch()
 
+        
         layout.addLayout(button_layout)
 
         search_widget.setLayout(layout)
@@ -222,6 +241,17 @@ def create_results_tab(self):
 
         results_widget.setLayout(layout)
         return results_widget
+    
+    def _search_collections(self):
+        log_info("RevEng.AI | Searching collections")
+        self.progress = create_progress_dialog(self, "RevEng.AI Search Collections", "Searching collections...")
+        search_term = self.search_input.text().strip()
+        self.search_collections_thread = SearchCollectionsThread(self.match_functions, self.bv, search_term)
+        self.search_collections_thread.finished.connect(self._on_search_collections_finished)
+        self.search_collections_thread.start()
+        
+        self.progress.show()
+        QCoreApplication.processEvents() 
 
     def search_collections(self):
         """Search for collections based on input"""

features/match_functions/search_collections_thread.py

@@ -0,0 +1,22 @@
+from PySide6.QtCore import QThread, Signal
+from binaryninja import log_info, BinaryView
+
+class SearchCollectionsThread(QThread):
+    finished = Signal(bool, str)  # Signal for success/failure and error message
+    
+    def __init__(self, match_functions, bv: BinaryView, search_term):
+        super().__init__()
+        self.match_functions = match_functions
+        self.bv = bv
+        self.search_term = search_term
+        log_info(f"RevEng.AI | Search term: {search_term}")
+
+    def run(self):
+        try:
+            success, message = self.match_functions.search_collections(self.bv, self.search_term)
+            if success:
+                self.finished.emit(True, message)
+            else:
+                self.finished.emit(False, message)
+        except Exception as e:
+            self.finished.emit(False, str(e))