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
+

+
+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