diff --git a/README.md b/README.md index 1200f1d..c422266 100644 --- a/README.md +++ b/README.md @@ -1,27 +1,154 @@ -# RevEng.AI Binary Ninja Plugin +

+Official Binary Ninja Plugin for RevEng.AI -A Binary Ninja plugin for integrating with the RevEng.AI platform. +### Features Supported -## Installation +This plugin brings the power of RevEng.AI directly into Binary Ninja. Here are the main features currently supported: -1. Ensure you have Python 3.9 or later installed -2. Install the required dependencies: - ```bash - pip install -r requirements.txt - ``` -3. Copy the `revengai_bn` directory to your Binary Ninja plugins directory: - - Linux: `~/.binaryninja/plugins/` - - Windows: `%APPDATA%\Binary Ninja\plugins\` - - macOS: `~/Library/Application Support/Binary Ninja/plugins/` +- **Configuration Panel**: Easily configure your API credentials and platform settings. +- **Choose a Source**: Select the uploaded binary whose analysis will be used for other features such as matching and renaming. +- **Process Binary**: Upload the currently loaded binary in Binary Ninja to RevEng.AI for analysis. +- **Auto-Unstrip**: Automatically restore stripped symbols in your binary using our AI engine. +- **Function Matching**: Compare and match functions from your current binary with those in your existing collections. +--- -## Requirements +## Installation & Running ๐Ÿš€ -- Binary Ninja 5.0 or later -- Python 3.9 or later -- Internet connection for API access -- RevEng.AI API key +### Step 1: Locate Binary Ninja Plugins Folder + +Locate your Binary Ninja user plugin directory: + +- **Tools > Plugins > Open Plugin Folder** from the Binary Ninja menu + +> ๐Ÿ’ก **Tip**: This opens the correct path regardless of OS or install type. + +Expected output locations: + - Linux: `~/.binaryninja/plugins/` + - Windows: `%APPDATA%\Binary Ninja\plugins\` + - macOS: `~/Library/Application Support/Binary Ninja/plugins/` + + +### Step 2: Download & Install the Plugin + +1. Visit the releases page: https://github.com/RevEngAI/reai-ida/releases +2. Download the latest release (look for the most recent version). +3. Extract the contents into the opened plugin folder. +4. Ensure the folder structure looks like this: + +``` +Example in Linux... +~/.binaryninja/plugins/ + โ””โ”€โ”€ revengai/ + โ””โ”€โ”€ [plugin files...] +``` + +> ๐Ÿ–ผ๏ธ *Insert screenshot of the correct plugin folder structure* + +### Step 3: Install Dependencies + +In your system terminal (not inside Binary Ninja), move to the directory with the 'requirements.txt' file and install required dependencies using: + +```bash +pip install -r requirements.txt +``` + +Or directly from within Binary Ninjaโ€™s built-in Python terminal: + +```python +import subprocess +subprocess.check_call(['pip', 'install', '-r', '/path/to/requirements.txt']) # Change to your path to requirements.txt +``` + +--- + +## Using the Plugin โš™๏ธ + +Once installed, youโ€™ll find `RevEng.AI` listed in the Binary Ninja plugins toolbar menu. + + + +Make sure to restart Binary Ninja completely after installation. +Then, check the Plugins menu โ€” the RevEng.AI plugin should be visible. +Finally, load a binary and explore the features described below. + +### 1. Configure the Plugin + +Select `Configuration` from the menu to set up your API key and host. + + + +Clicking "Continue" will validate your API key. + +--- + +### 2. Process a Binary + +Upload the currently loaded binary to RevEng.AI: + +- Select `RevEng.AI > Process Binary` + + + +Before starting the process, you can add a PDB file and debug information, assign custom tags for better tracking, choose which AI model you want to use, and decide whether to keep the analysis private (default) or make it publicly available. +The plugin will handle the upload and initiate the analysis. Once completed, an internal analysis ID is assigned. + +--- + +### 3. Choose Source Analysis + +If you have already processed your binary on the platform or if there are publicly available analyses, you can select one as your working source. + +- Select `RevEng.AI > Choose Source` + + + +This is required before using some features like function matching or auto unstrip. + +--- + +### 4. Auto Unstrip + +Bring back symbol names automatically: + +- Select `RevEng.AI > Auto Unstrip` + + + +Functions will be renamed with the most likely matching names from your configured collections. + +--- + +### 5. Match Functions + +Use function matching to identify similar functions in other binaries or collections: + +- Click `RevEng.AI > Match Functions` + + + +Matched functions are displayed based on the given confidence value. You can navigate or rename based on the results. + +--- + +## Troubleshooting + +- Only Binary Ninja 3.0+ is supported +- Python 3.9 or later is required +- Ensure your API key is valid and your analysis contains function-level information + +## Software Requirements + +This plugin relies on: + +- [reait](https://github.com/RevEngAI/reait) +- requests +- PySide6 ## License This plugin is released under the GPL-2.0 license. + +## Disclaimer + +Binary Ninja is a trademark of Vector 35. This project is not affiliated with or endorsed by Vector 35. diff --git a/features/auto_unstrip/auto_unstrip_thread.py b/features/auto_unstrip/auto_unstrip_thread.py deleted file mode 100644 index 3946b13..0000000 --- a/features/auto_unstrip/auto_unstrip_thread.py +++ /dev/null @@ -1,19 +0,0 @@ -from PySide6.QtCore import QThread, Signal - -class AutoUnstripThread(QThread): - finished = Signal(bool, str) # Signal for success/failure and error message - - def __init__(self, auto_unstrip, bv): - super().__init__() - self.auto_unstrip = auto_unstrip - self.bv = bv - - def run(self): - try: - success, message = self.auto_unstrip.auto_unstrip(self.bv) - if success: - self.finished.emit(True, message) - else: - self.finished.emit(False, message) - except Exception as e: - self.finished.emit(False, str(e)) \ No newline at end of file diff --git a/features/choose_source/analysis_load_thread.py b/features/choose_source/analysis_load_thread.py deleted file mode 100644 index 613a682..0000000 --- a/features/choose_source/analysis_load_thread.py +++ /dev/null @@ -1,21 +0,0 @@ -from PySide6.QtCore import QThread, Signal -from binaryninja import log_error - -class AnalysisLoadThread(QThread): - finished = Signal(list) - error = Signal(str) - - def __init__(self, choose_source, bv): - super().__init__() - self.choose_source = choose_source - self.bv = bv - - def run(self): - try: - analysis = self.choose_source.get_analysis(self.bv) - if not len(analysis): - raise Exception("No analysis found, try processing the binary again.") - self.finished.emit(analysis) - except Exception as e: - log_error(f"RevEng.AI | Failed to load analysis: {str(e)}") - self.error.emit(str(e)) \ No newline at end of file diff --git a/features/choose_source/choose_source_thread.py b/features/choose_source/choose_source_thread.py deleted file mode 100644 index 7afc077..0000000 --- a/features/choose_source/choose_source_thread.py +++ /dev/null @@ -1,20 +0,0 @@ -from PySide6.QtCore import QThread, Signal - -class ChooseSourceThread(QThread): - finished = Signal(bool, str) - - def __init__(self, choose_source, bv, chose): - super().__init__() - self.choose_source = choose_source - self.bv = bv - self.chose = chose - - def run(self): - try: - success = self.choose_source.choose_source(self.bv, self.chose) # Change to return error message - if success: - self.finished.emit(True, "") - else: - self.finished.emit(False, "") - except Exception as e: - self.finished.emit(False, str(e)) \ No newline at end of file diff --git a/features/upload/model_load_thread.py b/features/upload/model_load_thread.py deleted file mode 100644 index 1379adf..0000000 --- a/features/upload/model_load_thread.py +++ /dev/null @@ -1,19 +0,0 @@ -from PySide6.QtCore import QThread, Signal -from binaryninja import log_error - -class ModelLoadThread(QThread): - finished = Signal(list) - error = Signal(str) - - def __init__(self, uploader, bv): - super().__init__() - self.uploader = uploader - self.bv = bv - - def run(self): - try: - models = self.uploader.get_models(self.bv) - self.finished.emit(models) - except Exception as e: - log_error(f"RevEng.AI | Failed to load models: {str(e)}") - self.error.emit(str(e)) \ No newline at end of file diff --git a/features/upload/upload_thread.py b/features/upload/upload_thread.py deleted file mode 100644 index ea82219..0000000 --- a/features/upload/upload_thread.py +++ /dev/null @@ -1,20 +0,0 @@ -from PySide6.QtCore import QThread, Signal - -class UploadBinaryThread(QThread): - finished = Signal(bool, str) - - def __init__(self, uploader, bv, options): - super().__init__() - self.uploader = uploader - self.bv = bv - self.options = options - - def run(self): - try: - success = self.uploader.upload_binary(self.bv, self.options) # Change to return error message - if success: - self.finished.emit(True, "") - else: - self.finished.emit(False, "") - except Exception as e: - self.finished.emit(False, str(e)) \ No newline at end of file diff --git a/images/autounstrip.png b/images/autounstrip.png new file mode 100644 index 0000000..3f3833b Binary files /dev/null and b/images/autounstrip.png differ diff --git a/images/banner.png b/images/banner.png new file mode 100644 index 0000000..4befedf Binary files /dev/null and b/images/banner.png differ diff --git a/images/choosesource.png b/images/choosesource.png new file mode 100644 index 0000000..0904ac7 Binary files /dev/null and b/images/choosesource.png differ diff --git a/images/config.png b/images/config.png new file mode 100644 index 0000000..e6feeda Binary files /dev/null and b/images/config.png differ diff --git a/images/matchedfunctions.png b/images/matchedfunctions.png new file mode 100644 index 0000000..317f5ba Binary files /dev/null and b/images/matchedfunctions.png differ diff --git a/images/plugintoolbar.png b/images/plugintoolbar.png new file mode 100644 index 0000000..39041a0 Binary files /dev/null and b/images/plugintoolbar.png differ diff --git a/images/processbinary.png b/images/processbinary.png new file mode 100644 index 0000000..7aeab21 Binary files /dev/null and b/images/processbinary.png differ diff --git a/__init__.py b/revengai/__init__.py similarity index 100% rename from __init__.py rename to revengai/__init__.py diff --git a/features/__init__.py b/revengai/features/__init__.py similarity index 53% rename from features/__init__.py rename to revengai/features/__init__.py index d3c5c25..d5dcf85 100644 --- a/features/__init__.py +++ b/revengai/features/__init__.py @@ -2,5 +2,7 @@ from .upload import UploadFeature from .auto_unstrip import AutoUnstripFeature from .choose_source import ChooseSourceFeature +from .match_functions import MatchFunctionsFeature +from .match_current_function import MatchCurrentFunctionFeature -__all__ = ['ConfigurationFeature', 'UploadFeature', 'AutoUnstripFeature', 'ChooseSourceFeature'] \ No newline at end of file +__all__ = ['ConfigurationFeature', 'UploadFeature', 'AutoUnstripFeature', 'ChooseSourceFeature', 'MatchFunctionsFeature', 'MatchCurrentFunctionFeature'] \ No newline at end of file diff --git a/features/auto_unstrip/__init__.py b/revengai/features/auto_unstrip/__init__.py similarity index 90% rename from features/auto_unstrip/__init__.py rename to revengai/features/auto_unstrip/__init__.py index 97e5a8c..baceefd 100644 --- a/features/auto_unstrip/__init__.py +++ b/revengai/features/auto_unstrip/__init__.py @@ -1,7 +1,7 @@ from binaryninja import PluginCommand, log_info, BinaryView from .auto_unstrip import AutoUnstrip from .auto_unstrip_dialog import AutoUnstripDialog -from revengai_bn.utils import BaseAuthFeature +from revengai.utils import BaseAuthFeature class AutoUnstripFeature(BaseAuthFeature): def __init__(self, config=None): @@ -11,7 +11,7 @@ def __init__(self, config=None): def register(self): PluginCommand.register( - "RevEng.AI\\AutoUnstrip", + "RevEng.AI\\4 - Auto Unstrip", "Attempt to recover stripped function names", self.show_auto_unstrip_dialog, self.is_valid diff --git a/features/auto_unstrip/auto_unstrip.py b/revengai/features/auto_unstrip/auto_unstrip.py similarity index 77% rename from features/auto_unstrip/auto_unstrip.py rename to revengai/features/auto_unstrip/auto_unstrip.py index 2a9ac93..fe4438a 100644 --- a/features/auto_unstrip/auto_unstrip.py +++ b/revengai/features/auto_unstrip/auto_unstrip.py @@ -3,6 +3,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed from typing import List, Dict, Tuple import math +from revengai.utils import rename_function as rename_function_util class AutoUnstrip: def __init__(self, config): @@ -12,29 +13,7 @@ def __init__(self, config): self.path = None self.max_workers = 4 - def _rename_function(self, bv: BinaryView, addr: int, new_name: str, new_name_mangled: str) -> bool: - try: - 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 or func.name == new_name_mangled: - log_info(f"RevEng.AI | Function at {hex(addr)} already has name {func.name}") - return False - - new_symbol = Symbol(SymbolType.FunctionSymbol, addr, new_name) - bv.define_user_symbol(new_symbol) - - log_info(f"RevEng.AI | Renamed function at {hex(addr)} to {new_name}") - return True - - except Exception as e: - log_error(f"RevEng.AI | Error renaming function at {hex(addr)}: {str(e)}") - return False - def _process_batch(self, function_ids: List[int], id_to_addr: Dict[int, int], bv: BinaryView) -> Tuple[int, List[str]]: - """Process a batch of function IDs and return the number of renamed functions""" try: functions_by_distance = RE_nearest_symbols_batch( function_ids=function_ids, @@ -46,9 +25,9 @@ def _process_batch(self, function_ids: List[int], id_to_addr: Dict[int, int], bv functions = [] for function in functions_by_distance: functions.append({"function_id": function['origin_function_id'], "function_name": function['nearest_neighbor_function_name']}) - log_info(f"RevEng.AI | Functions by distance: {functions}") + #log_info(f"RevEng.AI | Functions by distance: {functions}") functions_by_score = RE_name_score(functions).json()["data"] - log_info(f"RevEng.AI | Functions by score: {functions_by_score}") + #log_info(f"RevEng.AI | Functions by score: {functions_by_score}") renamed_count = 0 errors = [] for result in functions_by_distance: @@ -69,16 +48,17 @@ def _process_batch(self, function_ids: List[int], id_to_addr: Dict[int, int], bv for function in functions_by_score: if function['function_id'] == func_id: if function['box_plot']["average"] < 0.9: - log_info(f"RevEng.AI | Function {function['function_id']} has a score of {function['box_plot']["average"]:.2f} for name {function['function_name']}, skipping") + log_info(f"RevEng.AI | Function {function['function_id']} has a score of {function['box_plot']['average']:.2f} for name {new_name_mangled}, skipping") break else: - log_info(f"RevEng.AI | Function {function['function_id']} has a score of {function['box_plot']["average"]:.2f} for name {function['function_name']}, renaming") - if self._rename_function(bv, func_addr, new_name, new_name_mangled): + log_info(f"RevEng.AI | Function {function['function_id']} has a score of {function['box_plot']['average']:.2f} for name {new_name_mangled}, renaming") + if rename_function_util(bv, func_addr, new_name_mangled): renamed_count += 1 break except Exception as e: + log_error(f"RevEng.AI | Error processing function {result['origin_function_id']}: {str(e)}") errors.append(str(e)) return renamed_count, errors @@ -146,4 +126,4 @@ def auto_unstrip(self, bv: BinaryView): except Exception as e: log_error(f"RevEng.AI | Error: {str(e)}") - return False, str(e) \ No newline at end of file + return False, str(e) diff --git a/features/auto_unstrip/auto_unstrip_dialog.py b/revengai/features/auto_unstrip/auto_unstrip_dialog.py similarity index 60% rename from features/auto_unstrip/auto_unstrip_dialog.py rename to revengai/features/auto_unstrip/auto_unstrip_dialog.py index 50444b0..914001b 100644 --- a/features/auto_unstrip/auto_unstrip_dialog.py +++ b/revengai/features/auto_unstrip/auto_unstrip_dialog.py @@ -1,14 +1,12 @@ -from binaryninja import BinaryView, PluginCommand, log_info, log_error +from binaryninja import log_error from PySide6.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, - QPushButton, QLabel, QCheckBox, - QGroupBox, QRadioButton, QSpacerItem, - QSizePolicy) + QPushButton, QLabel) from PySide6.QtCore import Qt -from PySide6.QtGui import QPixmap, QIcon +from PySide6.QtGui import QPixmap from PySide6.QtCore import QCoreApplication from PySide6.QtWidgets import QMessageBox -from revengai_bn.utils import create_progress_dialog -from .auto_unstrip_thread import AutoUnstripThread +from revengai.utils import create_progress_dialog +from revengai.utils.data_thread import DataThread import os class AutoUnstripDialog(QDialog): @@ -25,20 +23,28 @@ def init_ui(self): layout = QVBoxLayout() + header_layout = QHBoxLayout() + + logo_label = QLabel() + logo_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "images", "logo.png") + if os.path.exists(logo_path): + pixmap = QPixmap(logo_path) + pixmap = pixmap.scaled(100, 100, Qt.KeepAspectRatio, Qt.SmoothTransformation) + logo_label.setPixmap(pixmap) + header_layout.addWidget(logo_label) + info_layout = QVBoxLayout() title_label = QLabel("Auto Unstrip Binary") title_label.setStyleSheet("font-size: 18px; font-weight: bold;") - description_label = QLabel( - "Using official RevEng.AI sources, function names will be recovered based on a low similarity threshold and limited to available debug symbols.\nFunctions will be renamed automatically for easier analysis.\n\nThis process may take several minutes depending on the binary size." - ) + description_label = QLabel("Using official RevEng.AI sources, function names will be recovered based on a high similarity and confidence threshold and limited to available debug symbols.\nFunctions will be renamed automatically for easier analysis.\n\nThis process may take several minutes depending on the binary size.") description_label.setWordWrap(True) info_layout.addWidget(title_label) info_layout.addWidget(description_label) + header_layout.addLayout(info_layout, stretch=1) - layout.addLayout(info_layout) + layout.addLayout(header_layout) layout.addSpacing(20) - # Buttons button_layout = QHBoxLayout() self.save_button = QPushButton("Auto Unstrip") self.save_button.setStyleSheet(""" @@ -65,40 +71,24 @@ def init_ui(self): button_layout.addWidget(self.save_button) button_layout.addWidget(self.cancel_button) layout.addLayout(button_layout) - self.setLayout(layout) def _auto_unstrip(self): - log_info("RevEng.AI | Auto Unstripping binary") - # Create and show progress dialog using utility function self.progress = create_progress_dialog(self, "RevEng.AI Auto Unstrip", "Auto Unstripping binary...") + self.progress.show() + QCoreApplication.processEvents() - # Create and start upload thread - self.auto_unstrip_thread = AutoUnstripThread(self.auto_unstrip, self.bv) + self.auto_unstrip_thread = DataThread(self.auto_unstrip.auto_unstrip, self.bv) self.auto_unstrip_thread.finished.connect(self._on_auto_unstrip_finished) self.auto_unstrip_thread.start() - - self.progress.show() - QCoreApplication.processEvents() - def _on_auto_unstrip_finished(self, success, message): - """Handle auto unstrip completion""" self.progress.close() if success: - QMessageBox.information( - self, - "RevEng.AI Auto Unstrip", - f"Binary auto unstripped successfully!\n{message}", - QMessageBox.Ok - ) + QMessageBox.information(self, "RevEng.AI Auto Unstrip", f"Binary auto unstripped successfully!\n{message}", QMessageBox.Ok) self.accept() else: log_error(f"RevEng.AI | Failed to auto unstrip binary: {message}") - QMessageBox.critical( - self, - "RevEng.AI Auto Unstrip Error", - f"Failed to auto unstrip binary: {message}", - QMessageBox.Ok - ) \ No newline at end of file + QMessageBox.critical(self, "RevEng.AI Auto Unstrip Error", f"Failed to auto unstrip binary: {message}", QMessageBox.Ok) + self.reject() diff --git a/features/choose_source/__init__.py b/revengai/features/choose_source/__init__.py similarity index 90% rename from features/choose_source/__init__.py rename to revengai/features/choose_source/__init__.py index fcd852a..0fcfb8e 100644 --- a/features/choose_source/__init__.py +++ b/revengai/features/choose_source/__init__.py @@ -1,7 +1,7 @@ from binaryninja import PluginCommand, log_info, BinaryView from .choose_source import ChooseSource from .choose_source_dialog import ChooseSourceDialog -from revengai_bn.utils import BaseAuthFeature +from revengai.utils import BaseAuthFeature class ChooseSourceFeature(BaseAuthFeature): def __init__(self, config=None): @@ -11,7 +11,7 @@ def __init__(self, config=None): def register(self): PluginCommand.register( - "RevEng.AI\\Choose Source", + "RevEng.AI\\3 - Choose Source", "Choose a source for the binary analysis", self.show_choose_source_dialog, self.is_valid diff --git a/features/choose_source/choose_source.py b/revengai/features/choose_source/choose_source.py similarity index 81% rename from features/choose_source/choose_source.py rename to revengai/features/choose_source/choose_source.py index a449ba5..e498073 100644 --- a/features/choose_source/choose_source.py +++ b/revengai/features/choose_source/choose_source.py @@ -1,8 +1,5 @@ -from binaryninja import BinaryView, log_info, log_error, Symbol, SymbolType -from reait.api import RE_authentication, RE_search, RE_nearest_symbols_batch, RE_analyze_functions -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import List, Dict, Tuple -import math +from binaryninja import BinaryView, log_info, log_error +from reait.api import RE_search class ChooseSource: def __init__(self, config): @@ -15,15 +12,15 @@ def choose_source(self, bv: BinaryView, chose: str): binary_id = self.config.get_binary_id(bv) if binary_id == new_binary_id: log_info("RevEng.AI | Binary ID is already set to the chosen one.") - return True + return True, "Binary ID is already set to the chosen one." log_info(f"RevEng.AI | Changing Binary ID: {binary_id} to {new_binary_id}") self.config.set_current_info(new_binary_id) - return True + return True, "Binary ID changed successfully." except Exception as e: log_error(f"RevEng.AI | Failed to choose source: {str(e)}") - return False + return False, str(e) def get_analysis(self, bv: BinaryView): try: @@ -46,7 +43,7 @@ def get_analysis(self, bv: BinaryView): else: options.append(option) - return options + return True, options except Exception as e: log_error(f"RevEng.AI | Failed to get analysis: {str(e)}") - return [] \ No newline at end of file + return False, str(e) \ No newline at end of file diff --git a/features/choose_source/choose_source_dialog.py b/revengai/features/choose_source/choose_source_dialog.py similarity index 61% rename from features/choose_source/choose_source_dialog.py rename to revengai/features/choose_source/choose_source_dialog.py index 3851561..82e051b 100644 --- a/features/choose_source/choose_source_dialog.py +++ b/revengai/features/choose_source/choose_source_dialog.py @@ -1,11 +1,10 @@ from PySide6.QtWidgets import QDialog, QVBoxLayout, QLabel, QComboBox, QPushButton, QHBoxLayout from PySide6.QtGui import QPixmap from PySide6.QtCore import Qt, QCoreApplication -from .analysis_load_thread import AnalysisLoadThread -from .choose_source_thread import ChooseSourceThread +from revengai.utils.data_thread import DataThread from binaryninja import log_error from PySide6.QtWidgets import QMessageBox -from revengai_bn.utils import create_progress_dialog +from revengai.utils import create_progress_dialog import os class ChooseSourceDialog(QDialog): @@ -34,11 +33,7 @@ def init_ui(self): info_layout = QVBoxLayout() title_label = QLabel("Select Analysis Source") title_label.setStyleSheet("font-size: 18px; font-weight: bold;") - description_label = QLabel( - "Choose the source for your binary analysis. This selection will be used for all subsequent " - "features in the plugin, including auto-unstripping, function analysis, and other operations.\n\n" - #"The selected source will determine which database and models are used for your analysis tasks." - ) + description_label = QLabel("Choose the source for your binary analysis. This selection will be used for all subsequent features in the plugin, including auto-unstripping, function analysis, and other operations.\n\n") description_label.setWordWrap(True) info_layout.addWidget(title_label) info_layout.addWidget(description_label) @@ -87,65 +82,48 @@ def load_analysis(self): self.progress.show() QCoreApplication.processEvents() - self.analysis_thread = AnalysisLoadThread(self.choose_source, self.bv) + self.analysis_thread = DataThread(self.choose_source.get_analysis, self.bv) self.analysis_thread.finished.connect(self._on_analysis_loaded) - self.analysis_thread.error.connect(self._on_analysis_load_error) self.analysis_thread.start() - - def _on_analysis_loaded(self, analysis): + def _on_analysis_loaded(self, success, analysis): self.progress.close() self.combo.clear() - for analysis in analysis: - self.combo.addItem(analysis) - - def _on_analysis_load_error(self, error_msg): - self.progress.close() - log_error(f"RevEng.AI | Failed to load analysis: {error_msg}") - QMessageBox.critical( - self, - "RevEng.AI Analysis Loading Error", - f"Failed to load available analysis: {error_msg}", - QMessageBox.Ok - ) - self.reject() + if success: + if not len(analysis): + log_error("RevEng.AI | No analysis found, try processing the binary again.") + QMessageBox.critical(self, "RevEng.AI Analysis Loading Error", "No analysis found, try processing the binary again.", QMessageBox.Ok) + self.reject() + return + + for analysis in analysis: + self.combo.addItem(analysis) + else: + log_error(f"RevEng.AI | Failed to load analysis: {analysis}") + QMessageBox.critical(self, "RevEng.AI Analysis Loading Error", f"Failed to load available analysis: {analysis}", QMessageBox.Ok) + self.reject() def _choose_source(self): if not self.combo.currentText(): - log_warn("RevEng.AI | Source selection is required") - QMessageBox.warning( - self, - "RevEng.AI Choose Source", - "Please select a source for analysis.", - QMessageBox.Ok - ) + log_error("RevEng.AI | Source selection is required") + QMessageBox.warning(self, "RevEng.AI Choose Source", "Please select a source for analysis.", QMessageBox.Ok) return self.progress = create_progress_dialog(self, "RevEng.AI Choose Source", "Choosing source...") - - self.choose_source_thread = ChooseSourceThread(self.choose_source, self.bv, self.combo.currentText()) - self.choose_source_thread.finished.connect(self._on_choose_source_finished) - self.choose_source_thread.start() - self.progress.show() QCoreApplication.processEvents() + + self.choose_source_thread = DataThread(self.choose_source.choose_source, self.bv, self.combo.currentText()) + self.choose_source_thread.finished.connect(self._on_choose_source_finished) + self.choose_source_thread.start() - def _on_choose_source_finished(self, success, error_message): + def _on_choose_source_finished(self, success, message): self.progress.close() if success: - QMessageBox.information( - self, - "RevEng.AI Choose Source", - "Source chosen successfully!\nYou can now view the analysis on RevEng.AI", - QMessageBox.Ok - ) + QMessageBox.information(self, "RevEng.AI Choose Source", message, QMessageBox.Ok) self.accept() else: - log_error(f"RevEng.AI | Failed to choose source: {error_message}") - QMessageBox.critical( - self, - "RevEng.AI Choose Source Error", - f"Failed to choose source: {error_message}", - QMessageBox.Ok - ) \ No newline at end of file + log_error(f"RevEng.AI | Failed to choose source: {message}") + QMessageBox.critical(self, "RevEng.AI Choose Source Error", f"Failed to choose source: {message}", QMessageBox.Ok) + self.reject() \ No newline at end of file diff --git a/features/configuration/__init__.py b/revengai/features/configuration/__init__.py similarity index 86% rename from features/configuration/__init__.py rename to revengai/features/configuration/__init__.py index a9b02de..b86c64b 100644 --- a/features/configuration/__init__.py +++ b/revengai/features/configuration/__init__.py @@ -11,7 +11,7 @@ def __init__(self): def register(self): PluginCommand.register( - "RevEng.AI\\Configure", + "RevEng.AI\\1 - Configure", "Configure RevEng.AI settings", self.show_configuration ) @@ -26,8 +26,7 @@ def get_config(self): return self.config def _register_binary_view_event(self): - BinaryViewType.add_binaryview_finalized_event(self._add_binaryview_finalized_event) # TODO: Use binaryview_finalized_event instead, but without load 3 times - # TODO: Nao usar binaryview_finalized_event para checkar creds, resulta em comandos nao carregando.q + BinaryViewType.add_binaryview_finalized_event(self._add_binaryview_finalized_event) log_info("RevEng.AI | Registered binary view event handler") def _add_binaryview_finalized_event(self, bv): @@ -44,11 +43,11 @@ def _add_binaryview_finalized_event(self, bv): None, "RevEng.AI - Binary Not Found", "This binary has not been processed in the RevEng.AI platform yet.\n\n" - "Please upload and process the binary first using the 'RevEng.AI > Upload Binary' option " + "Please upload and process the binary first using the 'RevEng.AI > Process Binary' option " "before using other RevEng.AI features.", QMessageBox.Ok ) else: log_error(f"RevEng.AI | Configuration initialization failed: {message}") except Exception as e: - log_error(f"RevEng.AI | Error in binary view event handler: {str(e)}") \ No newline at end of file + log_error(f"RevEng.AI | Error in binary view event handler: {str(e)}") diff --git a/features/configuration/config.py b/revengai/features/configuration/config.py similarity index 100% rename from features/configuration/config.py rename to revengai/features/configuration/config.py diff --git a/features/configuration/config_dialog.py b/revengai/features/configuration/config_dialog.py similarity index 97% rename from features/configuration/config_dialog.py rename to revengai/features/configuration/config_dialog.py index 72a802f..708267a 100644 --- a/features/configuration/config_dialog.py +++ b/revengai/features/configuration/config_dialog.py @@ -5,7 +5,7 @@ from binaryninja import log_info, log_error, log_warn import os from .config_save_thread import ConfigSaveThread -from revengai_bn.utils import create_progress_dialog +from revengai.utils import create_progress_dialog class ConfigDialog(QDialog): def __init__(self, config): @@ -15,7 +15,6 @@ def __init__(self, config): self.progress = None self.init_ui() - def init_ui(self): self.setWindowTitle("RevEng.AI Configuration Wizard") self.setMinimumWidth(500) @@ -25,7 +24,7 @@ def init_ui(self): header_layout = QHBoxLayout() logo_label = QLabel() - logo_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "images", "logo.png") + logo_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "images", "logo.png") ## review that line if os.path.exists(logo_path): pixmap = QPixmap(logo_path) pixmap = pixmap.scaled(100, 100, Qt.KeepAspectRatio, Qt.SmoothTransformation) @@ -98,7 +97,6 @@ def init_ui(self): self.setLayout(layout) - def save_config(self): api_key = self.api_key_input.text().strip() host = self.host_input.text().strip() @@ -131,7 +129,6 @@ def save_config(self): self.progress.show() - def _on_save_finished(self, success, error_message): self.progress.close() diff --git a/features/configuration/config_save_thread.py b/revengai/features/configuration/config_save_thread.py similarity index 100% rename from features/configuration/config_save_thread.py rename to revengai/features/configuration/config_save_thread.py diff --git a/revengai/features/match_current_function/__init__.py b/revengai/features/match_current_function/__init__.py new file mode 100644 index 0000000..f234315 --- /dev/null +++ b/revengai/features/match_current_function/__init__.py @@ -0,0 +1,27 @@ +from binaryninja import PluginCommand, log_info, BinaryView +from .match_current_function import MatchCurrentFunction +from .match_current_function_dialog import MatchCurrentFunctionDialog +from revengai.utils import BaseAuthFeature + +class MatchCurrentFunctionFeature(BaseAuthFeature): + def __init__(self, config=None): + super().__init__(config) + self.match_current_function = MatchCurrentFunction(config) + log_info("RevEng.AI | MatchCurrentFunction Feature initialized") + + def register(self): + PluginCommand.register_for_address( + "RevEng.AI\\Match Current Function", + "Search and match the current function against RevEng.AI database", + self.show_match_current_function_dialog, + self.is_valid + ) + log_info("RevEng.AI | MatchCurrentFunction Feature registered") + + def show_match_current_function_dialog(self, bv: BinaryView, func): + log_info("RevEng.AI | Opening MatchCurrentFunction dialog") + dialog = MatchCurrentFunctionDialog(self.config, self.match_current_function, bv, func) + dialog.exec_() + + def is_valid(self, bv: BinaryView, func): + return self.config.is_configured == True \ No newline at end of file diff --git a/revengai/features/match_current_function/match_current_function.py b/revengai/features/match_current_function/match_current_function.py new file mode 100644 index 0000000..176b094 --- /dev/null +++ b/revengai/features/match_current_function/match_current_function.py @@ -0,0 +1,445 @@ +from binaryninja import BinaryView, log_info, log_error, Symbol, SymbolType +from reait.api import RE_authentication, RE_search, RE_nearest_symbols_batch, RE_analyze_functions, RE_collections_search, RE_binaries_search, RE_name_score, RE_functions_data_types, RE_functions_data_types_poll +from typing import List, Dict, Tuple, Optional, Any +from datetime import datetime +import os +import json +import re +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from revengai.utils import rename_function as rename_function_util + +class MatchCurrentFunction: + def __init__(self, config): + self.config = config + self.base_addr = None + self.path = None + self.binary_id = None + self.analyzed_functions = [] + self.filtered_collections = [] + self.filtered_binaries = [] + + def search_collections(self, bv: BinaryView, search_term: str = ""): + 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}") + if not self._is_query_empty(query): + items = self._search_collection(query) + log_info(f"RevEng.AI | Items: {items}") + return True, items + + except Exception as e: + log_error(f"RevEng.AI | Error searching collections: {str(e)}") + return False, str(e) + + def _process_batch(self, function_ids: List[int], id_to_addr: Dict[int, int], confidence_threshold: float, debug_symbols: bool, bv: BinaryView) -> Tuple[int, List[str]]: + """Process a batch of function IDs and return the number of matched functions and any errors""" + try: + log_info(f"RevEng.AI | Processing batch of {len(function_ids)} functions") + + functions_by_distance = RE_nearest_symbols_batch( + function_ids=function_ids, + debug_enabled=self.debug_symbols, + collections=self.filtered_collections, + binaries=self.filtered_binaries, + nns=self.debug_symbols_count + ).json()["function_matches"] + + functions = [] + for function in functions_by_distance: + functions.append({"function_id": function['origin_function_id'], "function_name": function['nearest_neighbor_function_name']}) + if len(functions) == 0: + return 0, [] + #log_info(f"RevEng.AI | Functions by distance: {functions}") + functions_by_score = RE_name_score(functions).json()["data"] + #log_info(f"RevEng.AI | Functions by score: {functions_by_score}") + matched_count = 0 + lines = [] + for result in functions_by_distance: + try: + + line = { + "icon_path": f"{os.path.dirname(__file__)}/../../images/failed.png", + "icon_text": "Failed", + "original_name": "N/A", + "matched_name": result['nearest_neighbor_function_name_mangled'] if result['nearest_neighbor_function_name_mangled'] else result['nearest_neighbor_function_name'], + "signature": "N/A", + "matched_binary": result['nearest_neighbor_binary_name'], + "similarity": f"{(result['confidence'] * 100):.2f}%", + "confidence": "N/A", + "error": "", + "nearest_neighbor_id": result['nearest_neighbor_id'], + "function_address": "N/A" + } + + func_addr = id_to_addr.get(result['origin_function_id']) + if not func_addr: + line["error"] = "Function not found in binary" + lines.append(line) + continue + + function = bv.get_function_at(func_addr) + if function: + line["original_name"] = function.name + line["function_address"] = function.start + + for function_by_score in functions_by_score: + if function_by_score['function_id'] == result['origin_function_id']: + + line["confidence"] = f"{function_by_score['box_plot']['average']:.2f}%" + + if not line["matched_name"] or line["matched_name"].startswith(("sub_", "FUN_")): + line["error"] = "Function name is also debug symbol" + log_info(f"RevEng.AI | Function name is also debug symbol: {line}") + break + + if function_by_score["box_plot"]["average"] < similarity_threshold: + line["error"] = "Function score is below confidence threshold" + break + else: + function = bv.get_function_at(id_to_addr.get(result['origin_function_id'])) + if not function: + log_error(f"RevEng.AI | Function not found: ID = {result['origin_function_id']} | Address = 0x{id_to_addr.get(result['origin_function_id']):x}") + line["icon_path"] = f"{os.path.dirname(__file__)}/../../images/success.png" + line["icon_text"] = "Success" + matched_count += 1 + break + + lines.append(line) + + except Exception as e: + log_error(f"RevEng.AI | Error processing function {result['origin_function_id']}: {str(e)}") + + return matched_count, lines + + except Exception as e: + log_error(f"RevEng.AI | Error processing batch: {str(e)}") + return 0, [str(e)] + + def match_functions(self, bv: BinaryView, options: Dict[str, Any]) -> List[Dict]: + """Match functions from the binary against RevEng.AI database""" + try: + log_info("RevEng.AI | Starting function matching") + + similarity_threshold = options.get("similarity_threshold", 90) * 0.01 + selected_collections = options.get("selected_collections", []) + debug_symbols = options.get("debug_symbols", False) + debug_symbols_count = options.get("debug_symbols_count", 5) + function_addr = options.get("function", None) + result = { "matched": 0, "skipped": 0, "data": [] } + + functions = bv.get_functions_containing(function_addr) + + if not functions: + log_error(f"RevEng.AI | Function not found at 0x{function_addr:x}") + raise Exception("Function not found at address") + + function = functions[0] + log_info(f"RevEng.AI | Function: {function.name} at 0x{function.start:x}") + + filtered_collections = [] + filtered_binaries = [] + for item in selected_collections: + if item["type"] == "Collection": + filtered_collections.append(item["id"]) + else: + filtered_binaries.append(item["id"]) + + log_info(f"RevEng.AI | Similarity threshold: {similarity_threshold}") + log_info(f"RevEng.AI | Selected collections: {selected_collections}") + log_info(f"RevEng.AI | Debug symbols: {debug_symbols}") + log_info(f"RevEng.AI | Debug symbols count: {debug_symbols_count}") + log_info(f"RevEng.AI | Function: 0x{function_addr:x}") + + binary_id = self.config.get_binary_id(bv) + if not binary_id: + raise Exception("Analysis not found. Please choose one using 'Choose Source' feature.") + + analyzed_functions = RE_analyze_functions(self.path, binary_id).json()["functions"] + + analyzed_function = next((f for f in analyzed_functions if (f["function_vaddr"] + bv.image_base) == function.start), None) + if not analyzed_function: + log_error(f"RevEng.AI | Function {function.name} not found in analyzed functions") + raise Exception("Function not found in analyzed functions") + + log_info(f"RevEng.AI | Found function {function.name} at 0x{function.start:x}") + + functions_by_distance = RE_nearest_symbols_batch( + function_ids=[analyzed_function["function_id"]], + distance=similarity_threshold, + debug_enabled=debug_symbols, + collections=filtered_collections, + binaries=filtered_binaries, + nns=debug_symbols_count + ).json()["function_matches"] + results = [] + for result in functions_by_distance: + try: + results.append({ + "original_name": function.name, + "matched_name": result['nearest_neighbor_function_name_mangled'] if result['nearest_neighbor_function_name_mangled'] else result['nearest_neighbor_function_name'], + "signature": "N/A", + "matched_binary": result['nearest_neighbor_binary_name'], + "similarity": f"{(result['confidence'] * 100):.2f}%", + "nearest_neighbor_id": result['nearest_neighbor_id'], + "function_address": function.start + }) + + except Exception as e: + log_error(f"RevEng.AI | Error processing function {result['origin_function_id']}: {str(e)}") + return True, results + + except Exception as e: + log_error(f"RevEng.AI | Error in function matching: {str(e)}") + raise + + def get_function_details(self, bv: BinaryView, function_address: int) -> Optional[Dict]: + """Get detailed information about a specific function""" + try: + function = bv.get_function_at(function_address) + if not function: + return None + + return { + "name": function.name, + "address": function.start, + "size": function.total_bytes, + "signature": str(function.type), + "basic_blocks": len(function.basic_blocks), + "call_sites": len(function.call_sites), + "callers": len(function.callers), + "callees": len(function.callees), + } + except Exception as e: + log_error(f"RevEng.AI | Error getting function details: {str(e)}") + return None + + def _parse_search_query(self, query: str) -> dict: + patterns = [ + "sha_256_hash", + "tag", + "binary_name", + "collection_name", + "function_name", + "model_name" + ] + + key_regex = "|".join(re.escape(p) for p in patterns) + regex = rf'\b({key_regex}):\s*([^:]+?)(?=,\s*(?:{key_regex}):|$)' + + matches = re.findall(regex, query) + + result = {key: None for key in patterns + ["query"]} + + for key, value in matches: + values = [v.strip() for v in value.split(',')] + result[key] = values if len(values) > 1 or key == "tag" else values[0] + + if not any(value is not None for value in result.values()): + result["query"] = query + + if result["tag"]: + result["tags"] = result["tag"] + del result["tag"] + + return result + + def _is_query_empty(self, query: dict) -> bool: + """Check if query is empty or contains only empty values""" + if not query: + return True + + return all(not str(v).strip() for v in query.values()) + + def _search_collection(self, query: Dict[str, Any] = {}) -> None: + + def parse_date(date_str: str) -> str: + dt = datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%S.%f") + return dt.strftime("%Y-%m-%d %H:%M:%S") + + def fetch_results(api_func, label: str) -> List[Dict[str, Any]]: + try: + log_info(f"RevEng.AI | Query: {query}") + response = api_func(query=query, page=1, page_size=1024).json() + results = response.get("data", {}).get("results", []) + log_info(f"Found {len(results)} {label.lower()}s") + return results + + except Exception as e: + log_error(f"RevEng.AI | Getting information failed. Reason: {str(e)}") + return [] + + def build_items(items_list: List[Dict[str, Any]], item_type: str) -> List[Tuple]: + items = [] + for item in items_list: + name_key = "collection_name" if item_type == "Collection" else "binary_name" + date_key = "last_updated_at" if item_type == "Collection" else "created_at" + id_key = "collection_id" if item_type == "Collection" else "binary_id" + icon = "lock.png" if item_type == "Collection" and item["scope"] == "PRIVATE" else \ + "unlock.png" if item_type == "Collection" else "file.png" + + items.append({ + "name": item[name_key], + "icon": icon, + "type": item_type, + "date": parse_date(item[date_key]), + "model_name": item["model_name"], + "owner": item["owned_by"], + "id": item[id_key] + }) + return items + + try: + + log_info(f"RevEng.AI | Searching for collections with '{query or 'N/A'}'") + + collections_data = fetch_results(RE_collections_search, "collection") + binaries_data = fetch_results(RE_binaries_search, "binary") + + table_items = build_items(collections_data, "Collection") + table_items += build_items(binaries_data, "Binary") + + return table_items + + except Exception as e: + log_error("Getting collections failed. Reason: %s", str(e)) + return False, str(e) + + def rename_function(self, bv: BinaryView, selected_result: Dict) -> List[Dict]: + try: + log_info(f"RevEng.AI | Starting function renaming for {len(selected_result)} functions") + + renamed_count = 0 + failed_count = 0 + + function_address = selected_result.get("function_address") + new_name = selected_result.get("matched_name") + + if not function_address or not new_name: + log_error(f"RevEng.AI | Missing function address or name for rename") + failed_count += 1 + return False, "Missing function address or name for rename" + + if rename_function_util(bv, function_address, new_name): + renamed_count += 1 + log_info(f"RevEng.AI | Successfully renamed function at {function_address:x} to {new_name}") + else: + failed_count += 1 + log_error(f"RevEng.AI | Failed to rename function at {function_address:x}") + + + message = f"Successfully renamed {renamed_count} functions" + if failed_count > 0: + message += f" ({failed_count} failed)" + + log_info(f"RevEng.AI | {message}") + return True, message + + except Exception as e: + log_error(f"RevEng.AI | Error in function renaming: {str(e)}") + return False, str(e) + + def _process_data_type_batch(self, chunk: List[Dict]) -> List[Dict]: + try: + log_info(f"RevEng.AI | Processing data types batch of {len(chunk)} items") + + nearest_neighbor_ids = [item["nearest_neighbor_id"] for item in chunk] + + response = RE_functions_data_types(nearest_neighbor_ids) + + if response.status_code != 200: + log_error(f"RevEng.AI | Data types API call failed with status {response.status_code}") + return [] + + data = response.json() + + if "status" in data and data["status"] == "processing": + poll_id = data.get("poll_id") + if poll_id: + log_info(f"RevEng.AI | Polling for data types with ID: {poll_id}") + + max_attempts = 30 + for attempt in range(max_attempts): + time.sleep(2) + poll_response = RE_functions_data_types_poll(poll_id) + + if poll_response.status_code == 200: + poll_data = poll_response.json() + if poll_data.get("status") == "completed": + data = poll_data + break + else: + log_error(f"RevEng.AI | Polling failed with status {poll_response.status_code}") + break + else: + log_error(f"RevEng.AI | Polling timed out after {max_attempts} attempts") + return [] + + signatures = [] + for item in data.get("data", []): + signature = self.make_signature(item.get("data_types", [])) + signatures.append({ + "nearest_neighbor_id": item["nearest_neighbor_id"], + "signature": signature + }) + + return signatures + + except Exception as e: + log_error(f"RevEng.AI | Error processing data types batch: {str(e)}") + return [] + + def make_signature(self, data_types: List[Dict]) -> str: + try: + if not data_types: + return "void function();" + + # For now, create a simple signature + # This would need to be enhanced based on actual data_types structure + return_type = "void" + params = [] + + for dt in data_types: + if dt.get("type") == "return": + return_type = dt.get("name", "void") + elif dt.get("type") == "parameter": + param_type = dt.get("name", "int") + param_name = dt.get("param_name", f"param{len(params)}") + params.append(f"{param_type} {param_name}") + + params_str = ", ".join(params) if params else "" + return f"{return_type} function({params_str});" + + except Exception as e: + log_error(f"RevEng.AI | Error creating signature: {str(e)}") + return "void function();" + + def fetch_data_types(self, bv: BinaryView, selected_results: List[Dict]) -> Tuple[bool, Dict[str, Any]]: + """Fetch data types for selected function matches""" + try: + log_info(f"RevEng.AI | Starting data type fetching for {len(selected_results)} functions") + + # Process in chunks to avoid API limits + chunk_size = 50 + all_signatures = [] + + for i in range(0, len(selected_results), chunk_size): + chunk = selected_results[i:i+chunk_size] + log_info(f"RevEng.AI | Processing chunk {i//chunk_size + 1}/{(len(selected_results) + chunk_size - 1)//chunk_size}") + + signatures = self._process_data_type_batch(chunk) + all_signatures.extend(signatures) + + success_count = len([s for s in all_signatures if s.get("signature") != "void function();"]) + + log_info(f"RevEng.AI | Data type fetching completed. {success_count} functions have signatures") + + return True, { + "signatures": all_signatures, + "success_count": success_count + } + + except Exception as e: + log_error(f"RevEng.AI | Error in data type fetching: {str(e)}") + return False, str(e) \ No newline at end of file diff --git a/revengai/features/match_current_function/match_current_function_dialog.py b/revengai/features/match_current_function/match_current_function_dialog.py new file mode 100644 index 0000000..6278b1f --- /dev/null +++ b/revengai/features/match_current_function/match_current_function_dialog.py @@ -0,0 +1,330 @@ +from binaryninja import BinaryView, log_info, log_error +from PySide6.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QPushButton, + QLabel, QLineEdit, QTableWidget, QTableWidgetItem, + QHeaderView, QTabWidget, QWidget, QMessageBox, + QCheckBox, QDoubleSpinBox, QSpinBox, QGroupBox, + QSplitter, QTextEdit, QProgressBar, QSlider) +from PySide6.QtCore import Qt, QTimer, QCoreApplication +from PySide6.QtGui import QIcon +from revengai.utils import create_progress_dialog +from revengai.utils.data_thread import DataThread +from .tab_search import SearchTab +from .tab_result import ResultTab + +class MatchCurrentFunctionDialog(QDialog): + def __init__(self, config, match_current_function, bv, func): + super().__init__() + self.config = config + self.match_current_function = match_current_function + self.bv = bv + self.func = func + self.init_ui() + + def init_ui(self): + self.setWindowTitle("RevEng.AI: Match Current Function") + self.setMinimumSize(1000, 700) + self.resize(1200, 800) + + main_layout = QVBoxLayout() + + self.tab_widget = QTabWidget() + + # Footer layout + footer_layout = self.create_footer_layout() + + # Search tab + self.search_tab = SearchTab(self.match_current_function, self.bv, self.status_label) + self.tab_widget.addTab(self.search_tab, "Search") + + # Results tab + self.results_tab = ResultTab(self.match_current_function, self.bv, self.status_label) + self.tab_widget.addTab(self.results_tab, "Results") + + main_layout.addWidget(self.tab_widget) + + main_layout.addLayout(footer_layout) + + # Status bar + #self.status_label = QLabel("Ready") + #main_layout.addWidget(self.status_label) + + self.setLayout(main_layout) + + def create_footer_layout(self): + footer_layout = QVBoxLayout() + + # Match settings section + settings_group = QGroupBox() + settings_layout = QVBoxLayout() + + # similarity slider + similarity_layout = QHBoxLayout() + similarity_layout.addWidget(QLabel("Similarity:")) + self.similaritySlider = QSlider() + self.similaritySlider.setMaximum(100) + self.similaritySlider.setPageStep(5) + self.similaritySlider.setSliderPosition(90) + self.similaritySlider.setOrientation(Qt.Horizontal) + self.similaritySlider.setInvertedAppearance(False) + self.similaritySlider.setInvertedControls(False) + self.similaritySlider.setTickPosition(QSlider.TicksBothSides) + self.similaritySlider.setTickInterval(5) + self.similaritySlider.setObjectName("similaritySlider") + similarity_layout.addWidget(self.similaritySlider) + + # Add similarity value label + self.similarity_value_label = QLabel("90") + self.similaritySlider.valueChanged.connect(lambda value: self.similarity_value_label.setText(str(value))) + similarity_layout.addWidget(self.similarity_value_label) + + settings_layout.addLayout(similarity_layout) + + # Debug symbols checkbox + self.debug_symbols_checkbox = QCheckBox("Limit Matches to Debug Symbols") + self.debug_symbols_checkbox.setChecked(True) + self.debug_symbols_spinbox = QSpinBox() + self.debug_symbols_spinbox.setMinimum(1) + self.debug_symbols_spinbox.setMaximum(20) + self.debug_symbols_spinbox.setValue(5) + self.debug_symbols_spinbox.setSuffix(" Functions") + self.debug_symbols_spinbox.setPrefix("Limit to ") + self.debug_symbols_spinbox.setFixedWidth(165) + + debug_symbols_layout = QHBoxLayout() + debug_symbols_layout.addWidget(self.debug_symbols_checkbox) + debug_symbols_layout.addWidget(self.debug_symbols_spinbox) + settings_layout.addLayout(debug_symbols_layout) + + settings_group.setLayout(settings_layout) + footer_layout.addWidget(settings_group) + + # Buttons layout + buttons_layout = QHBoxLayout() + + self.status_label = QLabel("Ready!") + buttons_layout.addWidget(self.status_label) + buttons_layout.addStretch() + + 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_fetching_data_types) + self.rename_selected_button.clicked.connect(self.start_renaming) + + for button in [self.fetch_data_types_button, self.rename_selected_button]: + button.setEnabled(False) + button.setStyleSheet(""" + QPushButton { + background-color: #474b4e; + color: white; + padding: 6px 12px; + border-radius: 4px; + } + """) + + buttons_layout.addWidget(self.fetch_results_button) + buttons_layout.addWidget(self.fetch_data_types_button) + buttons_layout.addWidget(self.rename_selected_button) + + footer_layout.addLayout(buttons_layout) + return footer_layout + + def start_matching(self): + """Start the current function matching process""" + similarity_threshold = self.similaritySlider.value() + + log_info("RevEng.AI | Starting current function matching process") + + # Create and show progress dialog + self.progress = create_progress_dialog(self, "RevEng.AI Match Current Function", "Matching current function...") + self.progress.show() + QCoreApplication.processEvents() + self.status_label.setText("Matching current function...") + + options = { + "similarity_threshold": similarity_threshold, + "selected_collections": self.search_tab.selected_collections, + "debug_symbols": self.debug_symbols_checkbox.isChecked(), + "debug_symbols_count": self.debug_symbols_spinbox.value(), + "function": self.func + } + + # Create and start matching thread + self.match_thread = DataThread( + self.match_current_function.match_functions, + self.bv, + options + ) + self.match_thread.finished.connect(self.on_matching_finished) + self.match_thread.start() + + def start_renaming(self): + """Start the current function renaming process""" + log_info("RevEng.AI | Starting current function renaming process") + + # Check if a result is selected + if not self.results_tab.selected_result: + QMessageBox.warning( + self, + "RevEng.AI Rename Function", + "Please select a function match first.", + QMessageBox.Ok + ) + return + + # Create and show progress dialog + self.progress = create_progress_dialog(self, "RevEng.AI Rename Selected Function", "Renaming selected function...") + self.progress.show() + QCoreApplication.processEvents() + self.status_label.setText("Renaming selected function...") + + # Create and start matching thread - pass single result as list for compatibility + self.rename_thread = DataThread( + self.match_current_function.rename_function, + self.bv, + self.results_tab.selected_result + ) + self.rename_thread.finished.connect(self.on_renaming_finished) + self.rename_thread.start() + + def start_fetching_data_types(self): + log_info("RevEng.AI | Starting current function data type fetching process") + try: + + # Create and show progress dialog + self.progress = create_progress_dialog(self, "RevEng.AI Fetch Data Types", "Fetching data types...") + self.progress.show() + QCoreApplication.processEvents() + self.status_label.setText("Fetching data types...") + + self.fetch_data_types_thread = DataThread( + self.match_current_function.fetch_data_types, + self.bv, + self.results_tab.current_matches + ) + self.fetch_data_types_thread.finished.connect(self.on_fetching_data_types_finished) + self.fetch_data_types_thread.start() + + log_info(f"RevEng.AI | Fetching data types thread started") + + except Exception as e: + log_error(f"RevEng.AI | Error starting data type fetching: {str(e)}") + if hasattr(self, 'progress'): + self.progress.close() + QMessageBox.critical( + self, + "RevEng.AI Fetch Data Types Error", + f"Failed to start data type fetching:\n{str(e)}", + QMessageBox.Ok + ) + + def on_renaming_finished(self, success, data): + """Handle renaming completion""" + self.progress.close() + if success: + log_info(f"RevEng.AI | Renaming completed: {data}") + QMessageBox.information( + self, + "RevEng.AI Rename Functions", + f"{data}", + QMessageBox.Ok + ) + else: + log_error(f"RevEng.AI | Renaming failed: {data}") + self.status_label.setText(f"Renaming failed: {data}") + QMessageBox.critical( + self, + "RevEng.AI Rename Functions Error", + QMessageBox.Ok + ) + + def on_matching_finished(self, success, data): + self.progress.close() + + if success: + self.results_tab.current_matches = data + self.results_tab.populate_results_table() + + if len(self.results_tab.current_matches) > 0: + for button in [self.fetch_data_types_button, self.rename_selected_button]: + button.setEnabled(True) + button.setStyleSheet(""" + QPushButton { + background-color: #6c757d; + color: white; + padding: 6px 12px; + border-radius: 4px; + } + QPushButton:hover { + background-color: #5a6268; + } + """) + + + # Switch to results tab + self.tab_widget.setCurrentIndex(1) + + #self.status_label.setText(f"Matching completed. Found {successful_count} successful matches.") + + QMessageBox.information( + self, + "RevEng.AI Match Current Function", + f"Current function matching completed successfully!\n" + f"Total functions found: {len(data)}", + QMessageBox.Ok + ) + else: + log_error(f"RevEng.AI | Current function matching failed: {data}") + self.status_label.setText(f"Matching failed: {data}") + QMessageBox.critical( + self, + "RevEng.AI Match Current Function Error", + f"Failed to match current function:\n{data}", + QMessageBox.Ok + ) + + def on_fetching_data_types_finished(self, success, data): + """Handle data type fetching completion""" + self.progress.close() + + if success: + log_info(f"RevEng.AI | Data type fetching completed with {data['success_count']} functions having signatures") + self.results_tab.update_current_matches_with_signatures(data["signatures"]) + self.results_tab.populate_results_table() + self.status_label.setText(f"Data type fetching completed: {data['success_count']} functions have signatures") + + QMessageBox.information( + self, + "RevEng.AI Fetch Data Types", + f"Data types fetched successfully.\n{data['success_count']} function{'' if data['success_count'] == 1 else 's'} have signatures.", + QMessageBox.Ok + ) + else: + log_error(f"RevEng.AI | Data type fetching failed: {data}") + self.status_label.setText(f"Data type fetching failed: {data}") + QMessageBox.critical( + self, + "RevEng.AI Fetch Data Types Error", + f"Failed to fetch data types:\n{data}", + QMessageBox.Ok + ) + + """ + def closeEvent(self, event): + self.accept() + """ \ No newline at end of file diff --git a/revengai/features/match_current_function/tab_result.py b/revengai/features/match_current_function/tab_result.py new file mode 100644 index 0000000..e62314b --- /dev/null +++ b/revengai/features/match_current_function/tab_result.py @@ -0,0 +1,147 @@ +from binaryninja import log_info, log_error +from PySide6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QPushButton, + QLabel, QLineEdit, QTableWidget, QTableWidgetItem, + QHeaderView, QGroupBox, QSlider, QCheckBox, QMessageBox) +from PySide6.QtCore import Qt, QCoreApplication +from PySide6.QtGui import QIcon +from revengai.utils import create_progress_dialog +from revengai.utils.data_thread import DataThread + +class ResultTab(QWidget): + + def __init__(self, match_current_function, bv, status_label): + super().__init__() + self.match_current_function = match_current_function + self.bv = bv + self.status_label = status_label + self.current_matches = [] + self.selected_result = {} + self.match_thread = None + + self.layout = QVBoxLayout() + self.setLayout(self.layout) + + self._build_result_section() + + def _build_result_section(self): + layout = QVBoxLayout() + + self.results_table = QTableWidget() + self.results_table.setColumnCount(6) + self.results_table.setHorizontalHeaderLabels([ + "Selected", "Original Function Name", "Matched Function Name", "Signature", "Matched Binary", "Similarity" + ]) + + header = self.results_table.horizontalHeader() + header.setSectionResizeMode(0, QHeaderView.ResizeToContents) + header.setSectionResizeMode(1, QHeaderView.ResizeToContents) + header.setSectionResizeMode(2, QHeaderView.ResizeToContents) + header.setSectionResizeMode(3, QHeaderView.ResizeToContents) + header.setSectionResizeMode(4, QHeaderView.ResizeToContents) + header.setSectionResizeMode(5, QHeaderView.ResizeToContents) + + self.results_table.setSelectionBehavior(QTableWidget.SelectRows) + self.results_table.setSelectionMode(QTableWidget.SingleSelection) + self.results_table.setAlternatingRowColors(True) + self.results_table.verticalHeader().setVisible(False) + + layout.addWidget(self.results_table) + + self.layout.addLayout(layout) + + def on_checkbox_changed(self, item_or_row, column=None): + """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() + is_checkbox = item_or_row.column() == 0 + else: # Called from cellClicked + row = item_or_row + is_checkbox = column == 0 + + # Get the match data for this row + if row < len(self.current_matches): + match = self.current_matches[row] + else: + return + + if match: + checkbox_item = self.results_table.item(row, 0) + current_state = checkbox_item.checkState() + + # Toggle state if clicked on non-checkbox cell + if is_checkbox: + new_state = current_state + else: + new_state = Qt.Unchecked if current_state == Qt.Checked else Qt.Checked + checkbox_item.setCheckState(new_state) + + # Handle unique selection (only one can be checked) + if new_state == Qt.Checked: + # Uncheck all other checkboxes + for i in range(self.results_table.rowCount()): + if i != row: + other_checkbox = self.results_table.item(i, 0) + if other_checkbox and other_checkbox.checkState() == Qt.Checked: + 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')}") + else: + # If unchecked, clear selection + self.selected_result = {} + log_info(f"RevEng.AI | Deselected function match") + + def populate_results_table(self): + """Populate the results table with function matches""" + self.selected_result = {} + self.results_table.setRowCount(len(self.current_matches)) + + # Safely disconnect existing connections + try: + self.results_table.itemChanged.disconnect() + except TypeError: + pass # No connections to disconnect + + for row, match in enumerate(self.current_matches): + select_item = QTableWidgetItem() + select_item.setFlags(Qt.ItemIsUserCheckable | Qt.ItemIsEnabled | Qt.ItemIsSelectable) + select_item.setCheckState(Qt.Unchecked) + self.results_table.setItem(row, 0, select_item) + + column_data = [ + "original_name", + "matched_name", + "signature", + "matched_binary", + "similarity", + ] + + # Create and set items for each column + for column, field in enumerate(column_data, start=1): + item = QTableWidgetItem(match.get(field, "N/A")) + item.setFlags(item.flags() & ~Qt.ItemIsEditable) + item.setData(Qt.UserRole, match) + item.setSelected(False) + self.results_table.setItem(row, column, item) + + # Connect signals after populating + self.results_table.itemChanged.connect(self.on_checkbox_changed) + + # Safely disconnect and reconnect cellClicked + try: + self.results_table.cellClicked.disconnect() + except TypeError: + pass # No connections to disconnect + + self.results_table.cellClicked.connect(self.on_checkbox_changed) + + def update_current_matches_with_signatures(self, signatures): + for match in self.current_matches: + if not match.get("nearest_neighbor_id"): + continue + for signature_data in signatures: + if match["nearest_neighbor_id"] == signature_data["nearest_neighbor_id"]: + match["signature"] = signature_data["signature"] + break + self.populate_results_table() \ No newline at end of file diff --git a/revengai/features/match_current_function/tab_search.py b/revengai/features/match_current_function/tab_search.py new file mode 100644 index 0000000..5cf94fb --- /dev/null +++ b/revengai/features/match_current_function/tab_search.py @@ -0,0 +1,185 @@ +from binaryninja import log_info, log_error +from PySide6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QPushButton, + QLabel, QLineEdit, QTableWidget, QTableWidgetItem, + QHeaderView, QGroupBox, QSlider, QCheckBox, QMessageBox) +from PySide6.QtCore import Qt, QCoreApplication +from PySide6.QtGui import QIcon +from revengai.utils import create_progress_dialog +from revengai.utils.data_thread import DataThread + +class SearchTab(QWidget): + + def __init__(self, match_current_function, bv, status_label): + super().__init__() + self.match_current_function = match_current_function + self.bv = bv + self.status_label = status_label + self.current_collections = [] + self.selected_collections = [] + self.search_collections_thread = None + + self.layout = QVBoxLayout() + self.setLayout(self.layout) + + self._build_search_section() + #self._build_settings_section() + + def _build_search_section(self): + search_group = QGroupBox() + search_layout = QVBoxLayout() + + search_input_layout = QHBoxLayout() + self.search_input = QLineEdit() + 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) + self.search_button.setStyleSheet(""" + QPushButton { + background-color: #007bff; + color: white; + padding: 6px 12px; + border-radius: 4px; + + } + QPushButton:hover { + background-color: #0056b3; + } + """) + + 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) + + self.collections_table = QTableWidget() + self.collections_table.setColumnCount(7) + self.collections_table.setHorizontalHeaderLabels([" ", "Name", "Type", "Date", "Model Name", "Owner", "ID"]) + + header = self.collections_table.horizontalHeader() + header.setSectionResizeMode(0, QHeaderView.ResizeToContents) + header.setSectionResizeMode(1, QHeaderView.Stretch) + header.setSectionResizeMode(2, QHeaderView.ResizeToContents) + header.setSectionResizeMode(3, QHeaderView.ResizeToContents) + header.setSectionResizeMode(4, QHeaderView.Stretch) + header.setSectionResizeMode(5, QHeaderView.ResizeToContents) + header.setSectionResizeMode(6, QHeaderView.Stretch) + + self.collections_table.setSelectionBehavior(QTableWidget.SelectRows) + self.collections_table.setSelectionMode(QTableWidget.MultiSelection) + self.collections_table.setAlternatingRowColors(True) + self.collections_table.verticalHeader().setVisible(False) + + search_layout.addWidget(self.collections_table) + search_group.setLayout(search_layout) + self.layout.addWidget(search_group) + + def _search_collections(self): + self.progress = create_progress_dialog(self, "RevEng.AI Search Collections", "Searching collections...") + self.progress.show() + QCoreApplication.processEvents() + + search_term = self.search_input.text().strip() + log_info(f"RevEng.AI | Search term: {search_term}") + + self.search_collections_thread = DataThread(self.match_current_function.search_collections, self.bv, search_term) + self.search_collections_thread.finished.connect(self._on_search_collections_finished) + self.search_collections_thread.start() + + def _on_search_collections_finished(self, success, data): + self.progress.close() + if success: + self.selected_collections.clear() + self.collections_table.clearSelection() + self.collections_table.setRowCount(0) + + self.current_collections = data + message = f"Found {len(self.current_collections)} collections!" + log_info(f"RevEng.AI | {message}") + self.status_label.setText(message) + self.populate_collections_table() + + else: + message = f"Error searching collections: {data}" + log_error(f"RevEng.AI | {message}") + self.status_label.setText(message) + QMessageBox.critical(self, "Search Error", message) + + def populate_collections_table(self): + self.collections_table.setRowCount(len(self.current_collections)) + self.collections_table.itemChanged.disconnect() + + for row, collection in enumerate(self.current_collections): + select_item = QTableWidgetItem() + select_item.setFlags(Qt.ItemIsUserCheckable | Qt.ItemIsEnabled | Qt.ItemIsSelectable) + select_item.setCheckState(Qt.Unchecked) + #select_item.setData(Qt.UserRole, collection) + self.collections_table.setItem(row, 0, select_item) + + columns = [ + (1, "name", lambda x: x), + (2, "type", lambda x: x), + (3, "date", lambda x: x), + (4, "model_name", lambda x: x), + (5, "owner", lambda x: x), + (6, "id", lambda x: str(x)) + ] + + for col_idx, field, transform in columns: + item = QTableWidgetItem(transform(collection.get(field, ""))) + item.setFlags(item.flags() & ~Qt.ItemIsEditable) + item.setData(Qt.UserRole, collection) + item.setSelected(False) + self.collections_table.setItem(row, col_idx, item) + + self.collections_table.itemChanged.connect(self.on_checkbox_changed) + try: + self.collections_table.cellClicked.disconnect() + except: + pass + self.collections_table.cellClicked.connect(self.on_checkbox_changed) + + def on_checkbox_changed(self, item_or_row, column=None): + if isinstance(item_or_row, QTableWidgetItem): # Called from itemChanged + row = item_or_row.row() + is_checkbox = item_or_row.column() == 0 + else: # Called from cellClicked + row = item_or_row + is_checkbox = column == 0 + + collection = self.collections_table.item(row, 1).data(Qt.UserRole) + collection_id = str(collection.get("id", "")) if collection else None + + if collection and collection_id: + checkbox_item = self.collections_table.item(row, 0) + current_state = checkbox_item.checkState() + + # Toggle state if clicked cell or checkbox changed + if is_checkbox: + new_state = current_state + else: + new_state = Qt.Unchecked if current_state == Qt.Checked else Qt.Checked + checkbox_item.setCheckState(new_state) + + # Update selected collections + if new_state == Qt.Checked: + if collection not in self.selected_collections: + self.selected_collections.append(collection) + else: + if collection in self.selected_collections: + self.selected_collections.remove(collection) + + log_info(f"RevEng.AI | Collection {'selected' if new_state == Qt.Checked else 'deselected'}: {collection.get('name', 'Unknown')}") + log_info(f"RevEng.AI | Selected collections: {len(self.selected_collections)}") + + # Update status + self.status_label.setText(f"Selected {len(self.selected_collections)} collections") + diff --git a/revengai/features/match_functions/__init__.py b/revengai/features/match_functions/__init__.py new file mode 100644 index 0000000..a3ebbb5 --- /dev/null +++ b/revengai/features/match_functions/__init__.py @@ -0,0 +1,24 @@ +from binaryninja import PluginCommand, log_info, BinaryView +from .match_functions import MatchFunctions +from .match_functions_dialog import MatchFunctionsDialog +from revengai.utils import BaseAuthFeature + +class MatchFunctionsFeature(BaseAuthFeature): + def __init__(self, config=None): + super().__init__(config) + self.match_functions = MatchFunctions(config) + log_info("RevEng.AI | MatchFunctions Feature initialized") + + def register(self): + PluginCommand.register( + "RevEng.AI\\5 - Match Functions", + "Search and match functions against RevEng.AI database", + self.show_match_functions_dialog, + self.is_valid + ) + log_info("RevEng.AI | MatchFunctions Feature registered") + + def show_match_functions_dialog(self, bv: BinaryView): + log_info("RevEng.AI | Opening MatchFunctions dialog") + dialog = MatchFunctionsDialog(self.config, self.match_functions, bv) + dialog.exec_() \ No newline at end of file diff --git a/revengai/features/match_functions/match_functions.py b/revengai/features/match_functions/match_functions.py new file mode 100644 index 0000000..fe2e61f --- /dev/null +++ b/revengai/features/match_functions/match_functions.py @@ -0,0 +1,452 @@ +from binaryninja import BinaryView, log_info, log_error +from reait.api import RE_nearest_symbols_batch, RE_analyze_functions, RE_collections_search, RE_binaries_search, RE_name_score, RE_functions_data_types, RE_functions_data_types_poll +from typing import List, Dict, Tuple, Optional, Any +from datetime import datetime +import os +import re +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from revengai.utils import rename_function as rename_function_util + +class MatchFunctions: + def __init__(self, config): + self.config = config + self.base_addr = None + self.path = None + self.binary_id = None + self.analyzed_functions = [] + self.filtered_collections = [] + self.filtered_binaries = [] + + def search_collections(self, bv: BinaryView, search_term: str = ""): + 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}") + if not self._is_query_empty(query): + items = self._search_collection(query) + log_info(f"RevEng.AI | Items: {items}") + return True, items + + except Exception as e: + log_error(f"RevEng.AI | Error searching collections: {str(e)}") + return False, str(e) + + def _process_batch(self, function_ids: List[int], id_to_addr: Dict[int, int], confidence_threshold: float, debug_symbols: bool, bv: BinaryView) -> Tuple[int, List[str]]: + try: + log_info(f"RevEng.AI | Processing batch of {len(function_ids)} functions") + + functions_by_distance = RE_nearest_symbols_batch( + function_ids=function_ids, + debug_enabled=debug_symbols, + collections=self.filtered_collections, + binaries=self.filtered_binaries, + nns=1 + ).json()["function_matches"] + + functions = [] + for function in functions_by_distance: + functions.append({"function_id": function['origin_function_id'], "function_name": function['nearest_neighbor_function_name']}) + if len(functions) == 0: + return 0, [] + functions_by_score = RE_name_score(functions).json()["data"] + matched_count = 0 + lines = [] + for result in functions_by_distance: + try: + + line = { + "icon_path": f"{os.path.dirname(__file__)}/../../images/failed.png", + "icon_text": "Failed", + "original_name": "N/A", + "matched_name": result['nearest_neighbor_function_name_mangled'] if result['nearest_neighbor_function_name_mangled'] else result['nearest_neighbor_function_name'], + "signature": "N/A", + "matched_binary": result['nearest_neighbor_binary_name'], + "similarity": f"{(result['confidence'] * 100):.2f}%", + "confidence": "N/A", + "error": "", + "nearest_neighbor_id": result['nearest_neighbor_id'], + "function_address": "N/A" + } + + func_addr = id_to_addr.get(result['origin_function_id']) + if not func_addr: + line["error"] = "Function not found in binary" + lines.append(line) + continue + + function = bv.get_function_at(func_addr) + if function: + line["original_name"] = function.name + line["function_address"] = function.start + + for function_by_score in functions_by_score: + if function_by_score['function_id'] == result['origin_function_id']: + + line["confidence"] = f"{function_by_score['box_plot']['average']:.2f}%" + + if not line["matched_name"] or line["matched_name"].startswith(("sub_", "FUN_")): + line["error"] = "Function name is also debug symbol" + log_info(f"RevEng.AI | Function name is also debug symbol: {line}") + break + + if function_by_score["box_plot"]["average"] < confidence_threshold: + line["error"] = "Function score is below confidence threshold" + break + else: + function = bv.get_function_at(id_to_addr.get(result['origin_function_id'])) + if not function: + log_error(f"RevEng.AI | Function not found: ID = {result['origin_function_id']} | Address = 0x{id_to_addr.get(result['origin_function_id']):x}") + line["icon_path"] = f"{os.path.dirname(__file__)}/../../images/success.png" + line["icon_text"] = "Success" + matched_count += 1 + break + + lines.append(line) + + except Exception as e: + log_error(f"RevEng.AI | Error processing function {result['origin_function_id']}: {str(e)}") + + return matched_count, lines + + except Exception as e: + log_error(f"RevEng.AI | Error processing batch: {str(e)}") + return 0, [str(e)] + + def match_functions(self, bv: BinaryView, options: Dict[str, Any]) -> List[Dict]: + try: + log_info("RevEng.AI | Starting function matching") + + confidence_threshold = options.get("confidence_threshold", 0.1) + selected_collections = options.get("selected_collections", []) + debug_symbols = options.get("debug_symbols", False) + result = { "matched": 0, "skipped": 0, "data": [] } + + self.filtered_collections = [] + self.filtered_binaries = [] + for item in selected_collections: + if item["type"] == "Collection": + self.filtered_collections.append(item["id"]) + else: + self.filtered_binaries.append(item["id"]) + + log_info(f"RevEng.AI | Confidence threshold: {confidence_threshold}") + log_info(f"RevEng.AI | Selected collections: {selected_collections}") + log_info(f"RevEng.AI | Debug symbols: {debug_symbols}") + + binary_id = self.config.get_binary_id(bv) + if not binary_id: + raise Exception("Analysis not found. Please choose one using 'Choose Source' feature.") + + analyzed_functions = RE_analyze_functions(self.path, binary_id).json()["functions"] + function_ids = [func["function_id"] for func in analyzed_functions] + + log_info(f"RevEng.AI | Found {len(function_ids)} functions to match") + + functions = bv.functions + len_functions = len(functions) + + log_info(f"RevEng.AI | Found {len_functions} functions and {len(analyzed_functions)} analyzed functions.") + + for index, function in enumerate(functions, 1): + #log_info( f"RevEng.AI | Searching for {function.name} [{index}/{len_functions}]") + + analyzed_function = next((f for f in analyzed_functions if (f["function_vaddr"] + bv.image_base) == function.start), None) + + if analyzed_function: + #log_info(f"RevEng.AI | Found function {function.name} at {function.start:x}") + function_ids.append(analyzed_function["function_id"]) + else: + result["skipped"] += 1 + result["data"].append({ + "icon_path": f"{os.path.dirname(__file__)}/../../images/failed.png", + "icon_text": "Failed", + "original_name": function.name, + "matched_name": "N/A", + "signature": "N/A", + "matched_binary": "N/A", + "similarity": "0.0%", + "confidence": "0.0%", + "error": "No Similar Function Found", + "function_address": function.start + }) + + chunk_size = 50 + chunks = [function_ids[i:i + chunk_size] for i in range(0, len(function_ids), chunk_size)] + + log_info(f"RevEng.AI | Processing {len(function_ids)} functions in {len(chunks)} chunks of size {chunk_size}") + + id_to_addr = { + func["function_id"]: func["function_vaddr"] + bv.image_base + for func in analyzed_functions + } + + total_matched_functions = 0 + with ThreadPoolExecutor(max_workers=4) as executor: + future_to_chunk = { + executor.submit(self._process_batch, chunk, id_to_addr, confidence_threshold, debug_symbols, bv): i + for i, chunk in enumerate(chunks) + } + + for future in as_completed(future_to_chunk): + chunk_index = future_to_chunk[future] + try: + matched_count, lines = future.result() + total_matched_functions += matched_count + result["data"].extend(lines) + log_info(f"RevEng.AI | Chunk {chunk_index} completed: matched {matched_count} functions") + except Exception as e: + log_error(f"RevEng.AI | Error processing chunk {chunk_index}: {str(e)}") + + result["matched"] = total_matched_functions + result["failed"] = len(analyzed_functions) - total_matched_functions - result["skipped"] + + def parse_confidence(item): + try: + return float(item["confidence"].strip('%')) + except (KeyError, ValueError): + return 0.0 + + sorted_list = sorted(result["data"], key=parse_confidence, reverse=True) + result["data"] = sorted_list + + return True, result + + except Exception as e: + log_error(f"RevEng.AI | Error matching functions: {str(e)}") + raise e + + def get_function_details(self, bv: BinaryView, function_address: int) -> Optional[Dict]: + try: + function = bv.get_function_at(function_address) + if not function: + return None + + return { + "name": function.name, + "address": hex(function_address), + "size": len(function), + "basic_blocks": len(function.basic_blocks), + "instructions": sum(len(bb) for bb in function.basic_blocks), + "call_sites": len(function.call_sites), + "callers": len(function.callers), + "callees": len(function.callees) + } + + except Exception as e: + log_error(f"RevEng.AI | Error getting function details: {str(e)}") + return None + + def _parse_search_query(self, query: str) -> dict: + patterns = [ + "sha_256_hash", + "tag", + "binary_name", + "collection_name", + "function_name", + "model_name" + ] + + key_regex = "|".join(re.escape(p) for p in patterns) + regex = rf'\b({key_regex}):\s*([^:]+?)(?=,\s*(?:{key_regex}):|$)' + + matches = re.findall(regex, query) + + result = {key: None for key in patterns + ["query"]} + + for key, value in matches: + values = [v.strip() for v in value.split(',')] + result[key] = values if len(values) > 1 or key == "tag" else values[0] + + if not any(value is not None for value in result.values()): + result["query"] = query + + if result["tag"]: + result["tags"] = result["tag"] + del result["tag"] + + return result + + def _is_query_empty(self, query: dict) -> bool: + """ + Check if the query dictionary is empty or contains only None values. + + Args: + query (dict): The query dictionary to check + + Returns: + bool: True if the query is empty, False otherwise + """ + return all(value is None for value in query.values()) + + def _search_collection(self, query: Dict[str, Any] = {}) -> None: + + def parse_date(date_str: str) -> str: + dt = datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%S.%f") + return dt.strftime("%Y-%m-%d %H:%M:%S") + + def fetch_results(api_func, label: str) -> List[Dict[str, Any]]: + try: + log_info(f"RevEng.AI | Query: {query}") + response = api_func(query=query, page=1, page_size=1024).json() + results = response.get("data", {}).get("results", []) + log_info(f"Found {len(results)} {label.lower()}s") + return results + + except Exception as e: + log_error(f"RevEng.AI | Getting information failed. Reason: {str(e)}") + return [] + + def build_items(items_list: List[Dict[str, Any]], item_type: str) -> List[Tuple]: + items = [] + for item in items_list: + name_key = "collection_name" if item_type == "Collection" else "binary_name" + date_key = "last_updated_at" if item_type == "Collection" else "created_at" + id_key = "collection_id" if item_type == "Collection" else "binary_id" + icon = "lock.png" if item_type == "Collection" and item["scope"] == "PRIVATE" else \ + "unlock.png" if item_type == "Collection" else "file.png" + + items.append({ + "name": item[name_key], + "icon": icon, + "type": item_type, + "date": parse_date(item[date_key]), + "model_name": item["model_name"], + "owner": item["owned_by"], + "id": item[id_key] + }) + return items + + try: + + log_info(f"RevEng.AI | Searching for collections with '{query or 'N/A'}'") + + collections_data = fetch_results(RE_collections_search, "collection") + binaries_data = fetch_results(RE_binaries_search, "binary") + + table_items = build_items(collections_data, "Collection") + table_items += build_items(binaries_data, "Binary") + + return table_items + + except Exception as e: + log_error("Getting collections failed. Reason: %s", str(e)) + return False, str(e) + + def rename_functions(self, bv: BinaryView, selected_results: List[Dict]) -> List[Dict]: + """Rename functions from the binary against RevEng.AI database""" + try: + log_info("RevEng.AI | Starting function renaming") + + renamed_count = 0 + for result in selected_results: + # Convert function_address from string to int + try: + addr = int(result['function_address']) + except (ValueError, TypeError): + log_error(f"RevEng.AI | Invalid function address: {result}") + continue + + if rename_function_util(bv, addr, result['matched_name']): + renamed_count += 1 + + success_message = f"Successfully renamed {renamed_count} functions!" if renamed_count > 0 else "No functions were renamed!" + + log_info(f"RevEng.AI | {success_message}") + + return True, success_message + except Exception as e: + log_error(f"RevEng.AI | Error renaming functions: {str(e)}") + return False, str(e) + + def _process_data_type_batch(self, chunk: List[Dict]) -> List[Dict]: + try: + log_info(f"RevEng.AI | Processing chunk of {len(chunk)} functions") + function_ids = set([result['nearest_neighbor_id'] for result in chunk]) + RE_functions_data_types(function_ids=list(function_ids)) + signatures = [] + while True: + response = RE_functions_data_types_poll( + function_ids=list(function_ids), + ).json() + data = response.get("data", {}) + total_count = data.get("total_count", 0) + total_data_types = data.get("total_data_types_count", 0) + items = data.get("items", []) + log_info(f"RevEng.AI | Response: {response}") + if total_count != total_data_types or all(item.get("completed", False) for item in items): + break + time.sleep(1) + + for item in items: + log_info(f"RevEng.AI | Item: {item['function_id']}") + for result in chunk: + if result['nearest_neighbor_id'] == item['function_id']: + signature = self.make_signature(item['data_types']) + if signature != "N/A": + signatures.append({"nearest_neighbor_id": result['nearest_neighbor_id'], "signature": signature}) + break + + log_info(f"RevEng.AI | Total count: {total_count}") + log_info(f"RevEng.AI | Total data types: {total_data_types}") + log_info(f"RevEng.AI | Items: {items}") + + return signatures + except Exception as e: + log_error(f"RevEng.AI | Error processing data type batch: {str(e)}") + return [] + + def make_signature(self, data_types: List[Dict]) -> str: + try: + log_info(f"RevEng.AI | Making signature for {data_types}") + signature = "" + signature += f"{data_types['func_types'].get('name', 'N/A')} " + + for dep in data_types['func_deps']: + signature += f"{dep.get('name', 'N/A')} " + + log_info(f"RevEng.AI | Signature: {signature}") + return signature + except Exception as e: + log_error(f"RevEng.AI | Error making signature: {str(e)}") + return "N/A" + + def fetch_data_types(self, bv: BinaryView, selected_results: List[Dict]) -> Tuple[bool, Dict[str, Any]]: + try: + log_info("RevEng.AI | Starting data type fetching") + + if len(selected_results) == 0: + return False, "No valid functions selected" + + chunk_size = 50 + chunks = [selected_results[i:i + chunk_size] for i in range(0, len(selected_results), chunk_size)] + + log_info(f"RevEng.AI | Processing {len(selected_results)} functions in {len(chunks)} chunks of size {chunk_size}") + + signatures = [] + + with ThreadPoolExecutor(max_workers=4) as executor: + future_to_chunk = { + executor.submit(self._process_data_type_batch, chunk): i + for i, chunk in enumerate(chunks) + } + + for future in as_completed(future_to_chunk): + chunk_index = future_to_chunk[future] + try: + chunk = future.result() + log_info(f"RevEng.AI | Chunk {chunk_index} completed") + signatures.extend(chunk) + + except Exception as e: + log_error(f"RevEng.AI | Error processing chunk {chunk_index}: {str(e)}") + + options = { + "success_count": len(signatures), + "signatures": signatures + } + + return True, options + except Exception as e: + log_error(f"RevEng.AI | Error fetching data types: {str(e)}") + return False, str(e) diff --git a/revengai/features/match_functions/match_functions_dialog.py b/revengai/features/match_functions/match_functions_dialog.py new file mode 100644 index 0000000..c7ea820 --- /dev/null +++ b/revengai/features/match_functions/match_functions_dialog.py @@ -0,0 +1,240 @@ +from binaryninja import log_info, log_error +from PySide6.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QPushButton, QLabel, QTabWidget, QMessageBox, QCheckBox, QGroupBox, QSlider) +from PySide6.QtCore import Qt, QCoreApplication +from revengai.utils import create_progress_dialog +from revengai.utils.data_thread import DataThread +from .tab_search import SearchTab +from .tab_result import ResultTab + +class MatchFunctionsDialog(QDialog): + def __init__(self, config, match_functions, bv): + super().__init__() + self.config = config + self.match_functions = match_functions + self.bv = bv + self.init_ui() + + def init_ui(self): + self.setWindowTitle("RevEng.AI: Match Functions") + self.setMinimumSize(1000, 700) + self.resize(1200, 800) + + main_layout = QVBoxLayout() + + self.tab_widget = QTabWidget() + footer_layout = self.create_footer_layout() + + self.search_tab = SearchTab(self.match_functions, self.bv, self.status_label) + self.tab_widget.addTab(self.search_tab, "Search") + + self.results_tab = ResultTab(self.match_functions, self.bv, self.status_label) + self.tab_widget.addTab(self.results_tab, "Results") + + main_layout.addWidget(self.tab_widget) + main_layout.addLayout(footer_layout) + + self.setLayout(main_layout) + + def create_footer_layout(self): + footer_layout = QVBoxLayout() + + settings_group = QGroupBox() + settings_layout = QVBoxLayout() + + 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) + + 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) + + self.debug_symbols_checkbox = QCheckBox("Limit Matches to Debug Symbols") + self.debug_symbols_checkbox.setChecked(True) + settings_layout.addWidget(self.debug_symbols_checkbox) + + settings_group.setLayout(settings_layout) + footer_layout.addWidget(settings_group) + + buttons_layout = QHBoxLayout() + + self.status_label = QLabel("Ready!") + buttons_layout.addWidget(self.status_label) + buttons_layout.addStretch() + + 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_fetching_data_types) + self.rename_selected_button.clicked.connect(self.start_renaming) + + for button in [self.fetch_data_types_button, self.rename_selected_button]: + button.setEnabled(False) + button.setStyleSheet(""" + QPushButton { + background-color: #474b4e; + color: white; + padding: 6px 12px; + border-radius: 4px; + } + """) + + buttons_layout.addWidget(self.fetch_results_button) + buttons_layout.addWidget(self.fetch_data_types_button) + buttons_layout.addWidget(self.rename_selected_button) + + footer_layout.addLayout(buttons_layout) + return footer_layout + + def start_matching(self): + confidence_threshold = self.confidenceSlider.value() + + log_info("RevEng.AI | Starting function matching process") + + self.progress = create_progress_dialog(self, "RevEng.AI Match Functions", "Matching functions...") + self.progress.show() + QCoreApplication.processEvents() + self.status_label.setText("Matching functions...") + + options = { + "confidence_threshold": confidence_threshold, + "selected_collections": self.search_tab.selected_collections, + "debug_symbols": self.debug_symbols_checkbox.isChecked() + } + + self.match_thread = DataThread(self.match_functions.match_functions, self.bv, options) + self.match_thread.finished.connect(self.on_matching_finished) + self.match_thread.start() + + def start_renaming(self): + log_info("RevEng.AI | Starting function renaming process") + + self.progress = create_progress_dialog(self, "RevEng.AI Rename Selected Functions", "Renaming Selected functions...") + self.progress.show() + QCoreApplication.processEvents() + self.status_label.setText("Renaming selected functions...") + + self.rename_thread = DataThread( + self.match_functions.rename_functions, + self.bv, + self.results_tab.selected_results + ) + self.rename_thread.finished.connect(self.on_renaming_finished) + self.rename_thread.start() + + def start_fetching_data_types(self): + log_info("RevEng.AI | Starting function data type fetching process") + + try: + self.progress = create_progress_dialog(self, "RevEng.AI Fetch Data Types", "Fetching data types...") + self.progress.show() + QCoreApplication.processEvents() + self.status_label.setText("Fetching data types...") + + if not hasattr(self.results_tab, 'selected_results') or not self.results_tab.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.results_tab.selected_results) + self.fetch_data_types_thread.finished.connect(self.on_fetching_data_types_finished) + self.fetch_data_types_thread.start() + + log_info(f"RevEng.AI | Fetching data types thread started") + + except Exception as e: + log_error(f"RevEng.AI | Error starting data type fetching: {str(e)}") + if hasattr(self, 'progress'): + self.progress.close() + QMessageBox.critical(self, "RevEng.AI Fetch Data Types Error", f"Failed to start data type fetching:\n{str(e)}", QMessageBox.Ok) + + def on_renaming_finished(self, success, data): + self.progress.close() + + if success: + log_info(f"RevEng.AI | Renaming completed: {data}") + QMessageBox.information(self, "RevEng.AI Rename Functions", f"{data}", QMessageBox.Ok) + else: + log_error(f"RevEng.AI | Renaming failed: {data}") + self.status_label.setText(f"Renaming failed: {data}") + QMessageBox.critical(self, "RevEng.AI Rename Functions Error", f"Failed to rename functions:\n{data}", QMessageBox.Ok) + + def on_matching_finished(self, success, data): + self.progress.close() + + if success: + self.results_tab.current_matches = data["data"] + self.results_tab.populate_results_table() + + if len(self.results_tab.selected_results) > 0: + for button in [self.fetch_data_types_button, self.rename_selected_button]: + button.setEnabled(True) + button.setStyleSheet(""" + QPushButton { + background-color: #6c757d; + color: white; + padding: 6px 12px; + border-radius: 4px; + } + QPushButton:hover { + background-color: #5a6268; + } + """) + + self.tab_widget.setCurrentIndex(1) + + successful_count = data["matched"] + skipped_count = data["skipped"] + failed_count = data["failed"] + total_count = successful_count + skipped_count + failed_count + + self.status_label.setText(f"Matching completed!") + self.results_tab.status_label.setText(f"Total Functions Analyzed: {total_count} | Successful Analyses: {successful_count} | Skipped Analyses: {skipped_count}") + QMessageBox.information(self, "RevEng.AI Match Functions", f"Function matching completed successfully!\nSuccessful matches: {successful_count}\nNot enough confidence: {failed_count}\nSkipped: {skipped_count}\nTotal functions analyzed: {total_count}", QMessageBox.Ok) + else: + log_error(f"RevEng.AI | Function matching failed: {data}") + self.status_label.setText(f"Matching failed: {data}") + QMessageBox.critical(self, "RevEng.AI Match Functions Error", f"Failed to match functions:\n{data}", QMessageBox.Ok) + + def on_fetching_data_types_finished(self, success, data): + self.progress.close() + + if success: + log_info(f"RevEng.AI | Data type fetching completed with {data['success_count']} functions having signatures") + self.results_tab.update_current_matches_with_signatures(data["signatures"]) + self.results_tab.populate_results_table() + self.status_label.setText(f"Data type fetching completed!") + self.results_tab.status_label.setText(f"Data type fetching completed: {data['success_count']} functions have signatures") + QMessageBox.information(self, "RevEng.AI Fetch Data Types", f"Data types fetched successfully.\n{data['success_count']} function{'' if data['success_count'] == 1 else 's'} have signatures.", QMessageBox.Ok) + else: + log_error(f"RevEng.AI | Data type fetching failed: {data}") + self.status_label.setText(f"Data type fetching failed: {data}") + QMessageBox.critical(self, "RevEng.AI Fetch Data Types Error", f"Failed to fetch data types:\n{data}", QMessageBox.Ok) \ No newline at end of file diff --git a/revengai/features/match_functions/tab_result.py b/revengai/features/match_functions/tab_result.py new file mode 100644 index 0000000..41ed30f --- /dev/null +++ b/revengai/features/match_functions/tab_result.py @@ -0,0 +1,112 @@ +from PySide6.QtWidgets import (QWidget, QVBoxLayout, QTableWidget, QTableWidgetItem, + QHeaderView, QAbstractItemView, QLineEdit, QLabel) +from PySide6.QtCore import Qt +from PySide6.QtGui import QIcon +from binaryninja import log_info + +class ResultTab(QWidget): + + def __init__(self, match_functions, bv, status_label): + super().__init__() + self.match_functions = match_functions + self.bv = bv + self.status_label = status_label + self.current_matches = [] + self.selected_results = [] + self.match_thread = None + + self.layout = QVBoxLayout() + self.setLayout(self.layout) + + self._build_result_section() + + def _build_result_section(self): + layout = QVBoxLayout() + + self.search_bar = QLineEdit() + self.search_bar.setPlaceholderText("Search results") + self.search_bar.textChanged.connect(self.filter_results) + layout.addWidget(self.search_bar) + + self.results_table = QTableWidget() + self.results_table.setColumnCount(8) + self.results_table.setHorizontalHeaderLabels([ + "Successful", "Original Function Name", "Matched Function Name", + "Signature", "Matched Binary", "Similarity", "Confidence", "Error" + ]) + self.results_table.setSelectionMode(QAbstractItemView.NoSelection) + + header = self.results_table.horizontalHeader() + header.setSectionResizeMode(0, QHeaderView.ResizeToContents) + header.setSectionResizeMode(1, QHeaderView.Stretch) + header.setSectionResizeMode(2, QHeaderView.Stretch) + header.setSectionResizeMode(3, QHeaderView.ResizeToContents) + header.setSectionResizeMode(4, QHeaderView.Stretch) + header.setSectionResizeMode(5, QHeaderView.ResizeToContents) + header.setSectionResizeMode(6, QHeaderView.ResizeToContents) + header.setSectionResizeMode(7, QHeaderView.Stretch) + + self.results_table.setAlternatingRowColors(True) + self.results_table.verticalHeader().setVisible(False) + #self.results_table.itemSelectionChanged.connect(self.on_result_selection_changed) + layout.addWidget(self.results_table) + + self.status_label = QLabel("No results yet") + layout.addWidget(self.status_label) + + self.layout.addLayout(layout) + + def filter_results(self, text): + log_info(f"RevEng.AI | Filtering results: {text}") + for row in range(self.results_table.rowCount()): + self.results_table.setRowHidden(row, True) + for col in range(self.results_table.columnCount()): + item = self.results_table.item(row, col) + if item: + if text.lower() in item.text().lower(): + log_info(f"RevEng.AI | Filtering results: {item.text()}") + self.results_table.setRowHidden(row, False) + break + + def populate_results_table(self): + self.selected_results.clear() + + self.results_table.setRowCount(len(self.current_matches)) + + for row, match in enumerate(self.current_matches): + icon_path = match.get("icon_path", "") + icon_text = match.get("icon_text", "") + icon_item = QTableWidgetItem() + icon_item.setIcon(QIcon(icon_path)) + icon_item.setText(icon_text) + icon_item.setFlags(icon_item.flags() & ~Qt.ItemIsEditable) + self.results_table.setItem(row, 0, icon_item) + + column_data = [ + "original_name", + "matched_name", + "signature", + "matched_binary", + "similarity", + "confidence", + "error" + ] + + for column, field in enumerate(column_data, start=1): + value = match.get(field, "N/A") + item = QTableWidgetItem(value if len(value) < 25 else value[:22] + "...") + item.setFlags(item.flags() & ~Qt.ItemIsEditable) + self.results_table.setItem(row, column, item) + + if icon_text == "Success": + self.selected_results.append(match) + + def update_current_matches_with_signatures(self, selected_results): + for match in self.current_matches: + if match.get("nearest_neighbor_id", False): + continue + for result in selected_results: + if match["nearest_neighbor_id"] == result["nearest_neighbor_id"]: + match["signature"] = result["signature"] + break + self.populate_results_table() \ No newline at end of file diff --git a/revengai/features/match_functions/tab_search.py b/revengai/features/match_functions/tab_search.py new file mode 100644 index 0000000..a44301a --- /dev/null +++ b/revengai/features/match_functions/tab_search.py @@ -0,0 +1,227 @@ +from binaryninja import log_info, log_error, log_debug +from PySide6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QPushButton, + QLabel, QLineEdit, QTableWidget, QTableWidgetItem, + QHeaderView, QGroupBox, QSlider, QCheckBox, QMessageBox) +from PySide6.QtCore import Qt, QCoreApplication +from revengai.utils import create_progress_dialog +from revengai.utils.data_thread import DataThread + +class SearchTab(QWidget): + + def __init__(self, match_functions, bv, status_label): + super().__init__() + self.match_functions = match_functions + self.bv = bv + self.status_label = status_label + self.current_collections = [] + self.selected_collections = [] + self.search_collections_thread = None + + self.layout = QVBoxLayout() + self.setLayout(self.layout) + + self._build_search_section() + #self._build_settings_section() + + def _build_search_section(self): + search_group = QGroupBox() + search_layout = QVBoxLayout() + + search_input_layout = QHBoxLayout() + self.search_input = QLineEdit() + 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) + self.search_button.setStyleSheet(""" + QPushButton { + background-color: #007bff; + color: white; + padding: 6px 12px; + border-radius: 4px; + + } + QPushButton:hover { + background-color: #0056b3; + } + """) + + 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) + + self.collections_table = QTableWidget() + self.collections_table.setColumnCount(7) + self.collections_table.setHorizontalHeaderLabels([" ", "Name", "Type", "Date", "Model Name", "Owner", "ID"]) + + header = self.collections_table.horizontalHeader() + header.setSectionResizeMode(0, QHeaderView.ResizeToContents) + header.setSectionResizeMode(1, QHeaderView.Stretch) + header.setSectionResizeMode(2, QHeaderView.ResizeToContents) + header.setSectionResizeMode(3, QHeaderView.ResizeToContents) + header.setSectionResizeMode(4, QHeaderView.Stretch) + header.setSectionResizeMode(5, QHeaderView.ResizeToContents) + header.setSectionResizeMode(6, QHeaderView.Stretch) + + self.collections_table.setSelectionBehavior(QTableWidget.SelectRows) + self.collections_table.setSelectionMode(QTableWidget.MultiSelection) + self.collections_table.setAlternatingRowColors(True) + self.collections_table.verticalHeader().setVisible(False) + + search_layout.addWidget(self.collections_table) + search_group.setLayout(search_layout) + self.layout.addWidget(search_group) + + def _build_settings_section(self): + settings_group = QGroupBox() + settings_layout = QVBoxLayout() + + 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) + + 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) + + self.debug_symbols_checkbox = QCheckBox("Limit Matches to Debug Symbols") + self.debug_symbols_checkbox.setChecked(True) + settings_layout.addWidget(self.debug_symbols_checkbox) + + settings_group.setLayout(settings_layout) + self.layout.addWidget(settings_group) + + def _search_collections(self): + self.progress = create_progress_dialog(self, "RevEng.AI Search Collections", "Searching collections...") + self.progress.show() + QCoreApplication.processEvents() + + search_term = self.search_input.text().strip() + log_info(f"RevEng.AI | Search term: {search_term}") + + self.search_collections_thread = DataThread(self.match_functions.search_collections, self.bv, search_term) + self.search_collections_thread.finished.connect(self._on_search_collections_finished) + self.search_collections_thread.start() + + def _on_search_collections_finished(self, success, data): + self.progress.close() + if success: + self.selected_collections.clear() + self.collections_table.clearSelection() + self.collections_table.setRowCount(0) + + self.current_collections = data + message = f"Found {len(self.current_collections)} collections!" + log_info(f"RevEng.AI | {message}") + self.status_label.setText(message) + self.populate_collections_table() + + else: + message = f"Error searching collections: {data}" + log_error(f"RevEng.AI | {message}") + self.status_label.setText(message) + QMessageBox.critical(self, "Search Error", message) + + def populate_collections_table(self): + self.collections_table.setRowCount(len(self.current_collections)) + self.collections_table.itemChanged.disconnect() + + for row, collection in enumerate(self.current_collections): + select_item = QTableWidgetItem() + select_item.setFlags(Qt.ItemIsUserCheckable | Qt.ItemIsEnabled | Qt.ItemIsSelectable) + select_item.setCheckState(Qt.Unchecked) + #select_item.setData(Qt.UserRole, collection) + self.collections_table.setItem(row, 0, select_item) + + columns = [ + (1, "name", lambda x: x), + (2, "type", lambda x: x), + (3, "date", lambda x: x), + (4, "model_name", lambda x: x), + (5, "owner", lambda x: x), + (6, "id", lambda x: str(x)) + ] + + for col_idx, field, transform in columns: + item = QTableWidgetItem(transform(collection.get(field, ""))) + item.setFlags(item.flags() & ~Qt.ItemIsEditable) + item.setData(Qt.UserRole, collection) + item.setSelected(False) + self.collections_table.setItem(row, col_idx, item) + + self.collections_table.itemChanged.connect(self.on_checkbox_changed) + try: + self.collections_table.cellClicked.disconnect() + except: + pass + self.collections_table.cellClicked.connect(self.on_checkbox_changed) + + def on_checkbox_changed(self, item_or_row, column=None): + if isinstance(item_or_row, QTableWidgetItem): # Called from itemChanged + row = item_or_row.row() + is_checkbox = item_or_row.column() == 0 + else: # Called from cellClicked + row = item_or_row + is_checkbox = column == 0 + + collection = self.collections_table.item(row, 1).data(Qt.UserRole) + collection_id = str(collection.get("id", "")) if collection else None + + if collection and collection_id: + checkbox_item = self.collections_table.item(row, 0) + current_state = checkbox_item.checkState() + + # Toggle state if clicked cell or checkbox changed + if is_checkbox: + new_state = current_state + else: + new_state = Qt.Unchecked if current_state == Qt.Checked else Qt.Checked + checkbox_item.setCheckState(new_state) + + is_selected = collection_id in [str(c.get("id", "")) for c in self.selected_collections] + + if new_state == Qt.Checked and not is_selected: + # Add to selection + self.selected_collections.append(collection) + # Select the row + for col in range(self.collections_table.columnCount()): + row_item = self.collections_table.item(row, col) + if row_item: + row_item.setSelected(True) + log_info(f"RevEng.AI | Added collection to selection: {collection.get('name', '')}") + + elif new_state == Qt.Unchecked and is_selected: + # Remove from selection + self.selected_collections = [c for c in self.selected_collections if str(c.get("id", "")) != collection_id] + # Deselect the row + for col in range(self.collections_table.columnCount()): + row_item = self.collections_table.item(row, col) + if row_item: + row_item.setSelected(False) + log_info(f"RevEng.AI | Removed collection from selection: {collection.get('name', '')}") + + # Update status + self.status_label.setText(f"Selected {len(self.selected_collections)} collection(s)") + log_info(f"RevEng.AI | Total selected collections: {len(self.selected_collections)}") + diff --git a/features/upload/__init__.py b/revengai/features/upload/__init__.py similarity index 90% rename from features/upload/__init__.py rename to revengai/features/upload/__init__.py index 43ddc6f..336084e 100644 --- a/features/upload/__init__.py +++ b/revengai/features/upload/__init__.py @@ -1,7 +1,7 @@ from binaryninja import PluginCommand, log_info, BinaryView from .upload import BinaryUploader from .upload_dialog import UploadDialog -from revengai_bn.utils import BaseAuthFeature +from revengai.utils import BaseAuthFeature class UploadFeature(BaseAuthFeature): def __init__(self, config=None): @@ -11,7 +11,7 @@ def __init__(self, config=None): def register(self): PluginCommand.register( - "RevEng.AI\\Process Binary", + "RevEng.AI\\2 - Process Binary", "Process current binary to RevEng.AI for analysis", self.show_upload_dialog, self.is_valid diff --git a/features/upload/upload.py b/revengai/features/upload/upload.py similarity index 60% rename from features/upload/upload.py rename to revengai/features/upload/upload.py index 87cdf9b..c888ec5 100644 --- a/features/upload/upload.py +++ b/revengai/features/upload/upload.py @@ -1,48 +1,20 @@ -from binaryninja import BinaryView, log_info, log_error, log_debug, SymbolType, BinaryViewType +from binaryninja import BinaryView, log_info, log_error from reait.api import RE_models, RE_upload, RE_analysis_lookup, RE_analyse -from revengai_bn.utils import PeriodicChecker +from revengai.utils import PeriodicChecker class BinaryUploader: def __init__(self, config): self.config = config - - + def get_models(self, bv: BinaryView): try: - guess_model_platform = "" - if bv.view_type == "PE": - guess_model_platform = "windows" - elif bv.view_type == "ELF": - guess_model_platform = "linux" - elif bv.view_type == "MACHO": - guess_model_platform = "macos" - - guess_model_arch = "" - if bv.arch.name == "x86": - guess_model_arch = "x86-32" - elif bv.arch.name == "x86_64": - guess_model_arch = "x86" - else: - guess_model_arch = bv.arch.name - - log_info(f"RevEng.AI | Architecture: {bv.arch.name} | File type: {bv.view_type}") - models = RE_models().json() - models = list([model["model_name"] for model in models["models"]]) - - guess_model = f"{guess_model_arch}-{guess_model_platform}" - log_info(f"RevEng.AI | Guess model: {guess_model}") - for i, model in enumerate(models): - if guess_model in model: - log_info(f"RevEng.AI | Found model: {model}") - models.insert(0, models.pop(i)) - break + models = RE_models().json()["data"]["models"] log_info(f"RevEng.AI | Models: {models}") - return models + return True, models except Exception as e: log_error(f"RevEng.AI | Failed to get models: {str(e)}") - return [] + return False, [] - def upload_binary(self, bv: BinaryView, options: dict): try: @@ -85,22 +57,18 @@ def upload_binary(self, bv: BinaryView, options: dict): skip_scraping=True, skip_sbom=True, skip_capabilities=True, - advanced_analysis=False + advanced_analysis=False, + duplicate=True ).json() - log_info(f"RevEng.AI | Analysis response: {analysis}") - analysis_info = RE_analysis_lookup(str(analysis["binary_id"])).json() - log_info(f"RevEng.AI | Binary ID: {analysis['binary_id']}") log_info(f"RevEng.AI | Analysis ID: {analysis_info['analysis_id']}") - - # TODO: Set binary and analysis id in config in form of id in array in settings - PeriodicChecker().start_checking(bv, analysis["binary_id"]) + PeriodicChecker().start_checking(bv, analysis["binary_id"], self.config.set_current_info) - return True + return True, "Analysis started successfully." except Exception as e: log_error(f"RevEng.AI | Failed to upload binary: {str(e)}") - return False \ No newline at end of file + return False, str(e) \ No newline at end of file diff --git a/features/upload/upload_dialog.py b/revengai/features/upload/upload_dialog.py similarity index 70% rename from features/upload/upload_dialog.py rename to revengai/features/upload/upload_dialog.py index e8904c8..a0607c7 100644 --- a/features/upload/upload_dialog.py +++ b/revengai/features/upload/upload_dialog.py @@ -1,11 +1,10 @@ from PySide6.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QLabel, QComboBox, QPushButton, QRadioButton, QButtonGroup, QLineEdit, QGroupBox, QFileDialog, QMessageBox) -from PySide6.QtCore import Qt, QCoreApplication -from binaryninja import log_info, log_error, log_warn -from .model_load_thread import ModelLoadThread -from .upload_thread import UploadBinaryThread -from revengai_bn.utils import create_progress_dialog +from PySide6.QtCore import QCoreApplication +from binaryninja import log_error +from revengai.utils import create_progress_dialog +from revengai.utils.data_thread import DataThread class UploadDialog(QDialog): def __init__(self, config, uploader, bv): @@ -86,82 +85,52 @@ def init_ui(self): self.load_models() - def load_models(self): self.progress = create_progress_dialog(self, "RevEng.AI", "Loading available models...") + self.progress.show() + QCoreApplication.processEvents() - self.model_thread = ModelLoadThread(self.uploader, self.bv) + self.model_thread = DataThread(self.uploader.get_models, self.bv) self.model_thread.finished.connect(self._on_models_loaded) - self.model_thread.error.connect(self._on_model_load_error) self.model_thread.start() - self.progress.show() - QCoreApplication.processEvents() - - def _on_models_loaded(self, models): + def _on_models_loaded(self, success, models): self.progress.close() self.model_combo.clear() - for model in models: - self.model_combo.addItem(model) + if success: + for model in models: + self.model_combo.addItem(model) + else: + log_error(f"RevEng.AI | Failed to load models: {models}") + QMessageBox.critical(self, "RevEng.AI Model Loading Error", f"Failed to load available models: {models}", QMessageBox.Ok) + self.reject() - def _on_model_load_error(self, error_msg): - self.progress.close() - log_error(f"RevEng.AI | Failed to load models: {error_msg}") - QMessageBox.critical( - self, - "RevEng.AI Model Loading Error", - f"Failed to load available models: {error_msg}", - QMessageBox.Ok - ) - self.reject() - def upload_binary(self): if not self.model_combo.currentText(): - log_warn("RevEng.AI | Model selection is required") - QMessageBox.warning( - self, - "RevEng.AI Upload", - "Please select a model for analysis.", - QMessageBox.Ok - ) + log_error("RevEng.AI | Model selection is required") + QMessageBox.warning(self, "RevEng.AI Upload", "Please select a model for analysis.", QMessageBox.Ok) return self.progress = create_progress_dialog(self, "RevEng.AI Upload", "Uploading binary to RevEng.AI...") + self.progress.show() + QCoreApplication.processEvents() - self.upload_thread = UploadBinaryThread(self.uploader, self.bv, self.get_upload_options()) + self.upload_thread = DataThread(self.uploader.upload_binary, self.bv, self.get_upload_options()) self.upload_thread.finished.connect(self._on_upload_finished) self.upload_thread.start() - self.progress.show() - QCoreApplication.processEvents() - def _on_upload_finished(self, success, error_message): self.progress.close() if success: - QMessageBox.information( - self, - "RevEng.AI Upload", - "Binary uploaded successfully!\nYou can now view the analysis on RevEng.AI", - QMessageBox.Ok - ) + QMessageBox.information(self, "RevEng.AI Upload", "Binary uploaded successfully!\nYou can now view the analysis on RevEng.AI", QMessageBox.Ok) self.accept() else: log_error(f"RevEng.AI | Failed to upload binary: {error_message}") - QMessageBox.critical( - self, - "RevEng.AI Upload Error", - f"Failed to upload binary: {error_message}", - QMessageBox.Ok - ) + QMessageBox.critical(self,"RevEng.AI Upload Error", f"Failed to upload binary: {error_message}", QMessageBox.Ok) def browse_debug_info(self): - file_path, _ = QFileDialog.getOpenFileName( - self, - "Select Debug Info or PDB", - "", - "Debug Info (*.pdb *.debug);;All Files (*.*)" - ) + file_path, _ = QFileDialog.getOpenFileName(self, "Select Debug Info or PDB", "", "Debug Info (*.pdb *.debug);;All Files (*.*)") if file_path: self.debug_combo.setCurrentText(file_path) diff --git a/revengai/images/failed.png b/revengai/images/failed.png new file mode 100644 index 0000000..096c2ad Binary files /dev/null and b/revengai/images/failed.png differ diff --git a/images/logo.png b/revengai/images/logo.png similarity index 100% rename from images/logo.png rename to revengai/images/logo.png diff --git a/revengai/images/success.png b/revengai/images/success.png new file mode 100644 index 0000000..de8fc0c Binary files /dev/null and b/revengai/images/success.png differ diff --git a/revengai.py b/revengai/revengai.py similarity index 69% rename from revengai.py rename to revengai/revengai.py index 999e1ef..2ab341d 100644 --- a/revengai.py +++ b/revengai/revengai.py @@ -3,6 +3,8 @@ from .features import UploadFeature from .features import AutoUnstripFeature from .features import ChooseSourceFeature +from .features import MatchFunctionsFeature +from .features import MatchCurrentFunctionFeature class RevengAIPlugin: def __init__(self): @@ -11,6 +13,8 @@ def __init__(self): self.upload_feature = UploadFeature(self.config_feature.get_config()) self.auto_unstrip_feature = AutoUnstripFeature(self.config_feature.get_config()) self.choose_source_feature = ChooseSourceFeature(self.config_feature.get_config()) + self.match_functions_feature = MatchFunctionsFeature(self.config_feature.get_config()) + self.match_current_function_feature = MatchCurrentFunctionFeature(self.config_feature.get_config()) self._register_features() def _register_features(self): @@ -19,4 +23,6 @@ def _register_features(self): self.upload_feature.register() self.auto_unstrip_feature.register() self.choose_source_feature.register() + self.match_functions_feature.register() + self.match_current_function_feature.register() \ No newline at end of file diff --git a/revengai/utils/__init__.py b/revengai/utils/__init__.py new file mode 100644 index 0000000..5f9bbad --- /dev/null +++ b/revengai/utils/__init__.py @@ -0,0 +1,6 @@ +from .periodic_check import PeriodicChecker +from .base_auth_feature import BaseAuthFeature +from .progress_dialog import create_progress_dialog, create_cancellable_progress_dialog +from .utils import rename_function, parse_date + +__all__ = ['PeriodicChecker', 'BaseAuthFeature', 'create_progress_dialog', 'create_cancellable_progress_dialog', 'rename_function', 'parse_date'] diff --git a/utils/base_auth_feature.py b/revengai/utils/base_auth_feature.py similarity index 100% rename from utils/base_auth_feature.py rename to revengai/utils/base_auth_feature.py diff --git a/revengai/utils/data_thread.py b/revengai/utils/data_thread.py new file mode 100644 index 0000000..e53dd34 --- /dev/null +++ b/revengai/utils/data_thread.py @@ -0,0 +1,27 @@ +from PySide6.QtCore import QThread, Signal +from binaryninja import log_info, BinaryView + +class DataThread(QThread): + finished = Signal(bool, object) + + def __init__(self, callback_function, bv: BinaryView, args = None): + super().__init__() + self.callback_function = callback_function + self.bv = bv + self.args = args + log_info(f"RevEng.AI | Data thread initialized") + + def run(self): + try: + if self.args is None: + success, content = self.callback_function(self.bv) + else: + success, content = self.callback_function(self.bv, self.args) + + if success: + log_info(f"RevEng.AI | Data thread finished with success") + self.finished.emit(True, content) + else: + self.finished.emit(False, content) + except Exception as e: + self.finished.emit(False, str(e)) \ No newline at end of file diff --git a/utils/periodic_check.py b/revengai/utils/periodic_check.py similarity index 84% rename from utils/periodic_check.py rename to revengai/utils/periodic_check.py index 552ed3c..e9c248a 100644 --- a/utils/periodic_check.py +++ b/revengai/utils/periodic_check.py @@ -4,6 +4,7 @@ from binaryninja import log_info, log_error, BinaryView from requests.exceptions import RequestException from reait.api import RE_status +from PySide6.QtWidgets import QMessageBox class PeriodicChecker: def __init__(self): @@ -15,7 +16,7 @@ def stop(self): self._current_timer = None log_info("RevEng.AI | Stopped periodic status check") - def start_checking(self, binary_view: BinaryView, binary_id: int, interval: float = 60) -> None: + def start_checking(self, binary_view: BinaryView, binary_id: int, config_callback, interval: float = 60) -> None: def _worker(bv: BinaryView, bid: int): try: response = RE_status(bv.file.filename, bid) @@ -38,7 +39,14 @@ def _worker(bv: BinaryView, bid: int): f"RevEng.AI | Scheduled next status check for: {basename(bv.file.filename)} [{bid}]" ) else: + config_callback(binary_id) log_info(f"RevEng.AI | Analysis completed with status: {status}") + QMessageBox.information( + None, + "RevEng.AI Analysis Complete", + f"Binary analysis completed!", + QMessageBox.Ok + ) except RequestException as ex: log_error(f"RevEng.AI | Error getting binary analysis status: {str(ex)}") diff --git a/revengai/utils/progress_dialog.py b/revengai/utils/progress_dialog.py new file mode 100644 index 0000000..421edac --- /dev/null +++ b/revengai/utils/progress_dialog.py @@ -0,0 +1,86 @@ +from PySide6.QtWidgets import QProgressDialog, QProgressBar, QPushButton +from PySide6.QtCore import Qt + +def create_progress_dialog(parent, title, message): + progress = QProgressDialog(message, None, 0, 0, parent) + progress.setWindowTitle(title) + progress.setWindowModality(Qt.WindowModal) + progress.setCancelButton(None) + progress.setMinimumWidth(400) + progress.setMinimumHeight(100) + + progress_bar = progress.findChild(QProgressBar) + if progress_bar: + progress_bar.setMinimumWidth(250) + progress_bar.setMinimumHeight(20) + + progress.setStyleSheet(""" + QProgressBar { + border: 1px solid #cccccc; + border-radius: 4px; + text-align: center; + background-color: #f0f0f0; + min-width: 250px; + min-height: 20px; + } + QProgressBar::chunk { + background-color: #007bff; + border-radius: 3px; + } + """) + + return progress + +def create_cancellable_progress_dialog(parent, title, message, cancel_callback=None): + """Create a progress dialog with a cancel button that can stop threads""" + progress = QProgressDialog(message, "Cancel", 0, 0, parent) + progress.setWindowTitle(title) + progress.setWindowModality(Qt.WindowModal) + progress.setMinimumWidth(400) + progress.setMinimumHeight(100) + + # Style the cancel button + cancel_button = progress.findChild(QPushButton) + if cancel_button: + cancel_button.setStyleSheet(""" + QPushButton { + background-color: #dc3545; + color: white; + padding: 6px 12px; + border-radius: 4px; + border: none; + min-width: 60px; + } + QPushButton:hover { + background-color: #c82333; + } + QPushButton:pressed { + background-color: #bd2130; + } + """) + + progress_bar = progress.findChild(QProgressBar) + if progress_bar: + progress_bar.setMinimumWidth(250) + progress_bar.setMinimumHeight(20) + + progress.setStyleSheet(""" + QProgressBar { + border: 1px solid #cccccc; + border-radius: 4px; + text-align: center; + background-color: #f0f0f0; + min-width: 250px; + min-height: 20px; + } + QProgressBar::chunk { + background-color: #007bff; + border-radius: 3px; + } + """) + + # Connect cancel callback if provided + if cancel_callback: + progress.canceled.connect(cancel_callback) + + return progress \ No newline at end of file diff --git a/revengai/utils/utils.py b/revengai/utils/utils.py new file mode 100644 index 0000000..6588cff --- /dev/null +++ b/revengai/utils/utils.py @@ -0,0 +1,30 @@ +from datetime import datetime +from binaryninja import BinaryView, log_error, log_info, Symbol, SymbolType + +def rename_function(bv: BinaryView, addr: int, new_name: str, data_type: dict = None) -> bool: + try: + 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 + + new_symbol = Symbol(SymbolType.FunctionSymbol, addr, new_name) + bv.define_user_symbol(new_symbol) + + log_info(f"RevEng.AI | Renamed function at {hex(addr)} to {new_name}") + return True + + except Exception as e: + log_error(f"RevEng.AI | Error renaming function at {hex(addr)}: {str(e)}") + return False + +def parse_date(date_str: str) -> str: + try: + dt = datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%S.%f") + return dt.strftime("%Y-%m-%d %H:%M:%S") + except Exception as e: + return date_str \ No newline at end of file diff --git a/utils/__init__.py b/utils/__init__.py deleted file mode 100644 index 72bde14..0000000 --- a/utils/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .periodic_check import PeriodicChecker -from .base_auth_feature import BaseAuthFeature -from .progress_dialog import create_progress_dialog - -__all__ = ['PeriodicChecker', 'BaseAuthFeature', 'create_progress_dialog'] diff --git a/utils/__pycache__/__init__.cpython-38.pyc b/utils/__pycache__/__init__.cpython-38.pyc deleted file mode 100644 index 962b0bc..0000000 Binary files a/utils/__pycache__/__init__.cpython-38.pyc and /dev/null differ diff --git a/utils/__pycache__/config_validator.cpython-38.pyc b/utils/__pycache__/config_validator.cpython-38.pyc deleted file mode 100644 index 13c94be..0000000 Binary files a/utils/__pycache__/config_validator.cpython-38.pyc and /dev/null differ diff --git a/utils/progress_dialog.py b/utils/progress_dialog.py deleted file mode 100644 index d355bea..0000000 --- a/utils/progress_dialog.py +++ /dev/null @@ -1,32 +0,0 @@ -from PySide6.QtWidgets import QProgressDialog, QProgressBar -from PySide6.QtCore import Qt - -def create_progress_dialog(parent, title, message): - progress = QProgressDialog(message, None, 0, 0, parent) - progress.setWindowTitle(title) - progress.setWindowModality(Qt.WindowModal) - progress.setCancelButton(None) - progress.setMinimumWidth(400) - progress.setMinimumHeight(100) - - progress_bar = progress.findChild(QProgressBar) - if progress_bar: - progress_bar.setMinimumWidth(250) - progress_bar.setMinimumHeight(20) - - progress.setStyleSheet(""" - QProgressBar { - border: 1px solid #cccccc; - border-radius: 4px; - text-align: center; - background-color: #f0f0f0; - min-width: 250px; - min-height: 20px; - } - QProgressBar::chunk { - background-color: #007bff; - border-radius: 3px; - } - """) - - return progress \ No newline at end of file