diff --git a/.revengai/features.json b/.revengai/features.json index ee62d78..0c1a9f1 100644 --- a/.revengai/features.json +++ b/.revengai/features.json @@ -9,9 +9,6 @@ "apply_existing_analysis": { "status": "yes" }, - "binary_auto_analysis": { - "status": "yes" - }, "binary_upload": { "note": "No explicit command or feature to upload binaries. Binaries are uploaded on creation of a new analysis only.", "status": "yes" diff --git a/reai_toolkit/features/__init__.py b/reai_toolkit/features/__init__.py index cc2e1f5..65ac54f 100755 --- a/reai_toolkit/features/__init__.py +++ b/reai_toolkit/features/__init__.py @@ -1,6 +1,5 @@ from reai_toolkit.features.configuration import ConfigurationFeature from reai_toolkit.features.upload import UploadFeature -from reai_toolkit.features.auto_unstrip import AutoUnstripFeature from reai_toolkit.features.choose_source import ChooseSourceFeature from reai_toolkit.features.match_functions import MatchFunctionsFeature from reai_toolkit.features.match_current_function import MatchCurrentFunctionFeature @@ -10,9 +9,8 @@ from reai_toolkit.features.view_analysis import ViewAnalysisFeature __all__ = [ 'ConfigurationFeature', - 'UploadFeature', - 'AutoUnstripFeature', - 'ChooseSourceFeature', + 'UploadFeature', + 'ChooseSourceFeature', 'MatchFunctionsFeature', 'MatchCurrentFunctionFeature', 'ViewFunctionInPortalFeature', diff --git a/reai_toolkit/features/auto_unstrip/__init__.py b/reai_toolkit/features/auto_unstrip/__init__.py deleted file mode 100755 index 67bd5b1..0000000 --- a/reai_toolkit/features/auto_unstrip/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -from reai_toolkit.features.auto_unstrip.auto_unstrip import AutoUnstrip -from reai_toolkit.utils import BaseAuthFeature -from reai_toolkit.features.auto_unstrip.auto_unstrip_dialog import AutoUnstripDialog -from binaryninja import PluginCommand, log_info, BinaryView - -class AutoUnstripFeature(BaseAuthFeature): - def __init__(self, config=None): - super().__init__(config) - self.auto_unstrip = AutoUnstrip(config) - log_info("RevEng.AI | AutoUnstrip Feature initialized") - - def register(self): - PluginCommand.register( - "RevEng.AI\\​Auto Unstrip", - "Attempt to recover stripped function names", - self.show_auto_unstrip_dialog, - self.is_valid - ) - log_info("RevEng.AI | AutoUnstrip Feature registered") - - def show_auto_unstrip_dialog(self, bv: BinaryView): - log_info("RevEng.AI | Opening AutoUnstrip dialog") - dialog = AutoUnstripDialog(self.config, self.auto_unstrip, bv) - dialog.exec_() - - def is_valid(self, bv: BinaryView): - return self.config.is_configured == True and self.config.analysis_id is not None - diff --git a/reai_toolkit/features/auto_unstrip/auto_unstrip.py b/reai_toolkit/features/auto_unstrip/auto_unstrip.py deleted file mode 100755 index a503b0d..0000000 --- a/reai_toolkit/features/auto_unstrip/auto_unstrip.py +++ /dev/null @@ -1,162 +0,0 @@ -import time -import revengai -from typing import List, Dict, Tuple -from libbs.api import DecompilerInterface -from libbs.artifacts import _art_from_dict, Function -from binaryninja import BinaryView, log_info, log_error -from libbs.decompilers.binja.interface import BinjaInterface -from concurrent.futures import ThreadPoolExecutor, as_completed -from reai_toolkit.utils import rename_function as rename_function_util, apply_data_types as apply_data_types_util, get_function_by_addr as get_function_by_addr_util -from revengai import AutoUnstripRequest -from threading import Event - -class AutoUnstrip: - def __init__(self, config): - self.config = config - self.auto_unstrip_distance = 0.09999999999999998 - self.base_addr = None - self.path = None - self.max_workers = 4 - self.cancelled = Event() - - def cancel(self): - log_info("RevEng.AI | Cancelling operation...") - self.cancelled.set() - - def clear_cancelled(self): - log_info("RevEng.AI | Clearing cancelled event...") - self.cancelled.clear() - - def resolve_data_types(self, to_datatypes: List[Dict], id_to_addr: Dict[int, int], deci: DecompilerInterface, chunk_index: int) -> None: - try: - function_ids = set([result['nearest_neighbor_id'] for result in to_datatypes]) - log_info(f"RevEng.AI | Resolving data types for {len(function_ids)} functions") - RE_functions_data_types(function_ids=list(function_ids)) - - items = [] - while True: - response = RE_functions_data_types_poll( - function_ids=list(function_ids), - ).json() - data = response.get("data", {}) - items = data.get("items", []) - pending_count = sum(1 for item in items if item.get("status") == "pending") - log_info(f"RevEng.AI | [Chunk {chunk_index}] {pending_count} items still pending...") - if not pending_count: - break - time.sleep(3) - - for item in items: - log_info(f"RevEng.AI | Item: {item['function_id']}") - if item['status'] != "completed": - continue - for result in to_datatypes: - if result['nearest_neighbor_id'] == item['function_id']: - - - signature = "N/A" - item2 = item.get("data_types", {}) - func_types = item2.get("func_types", None) - func_deps = item2.get("func_deps", []) - log_info(f"RevEng.AI | Func types: {func_types}") - if func_types is not None: - fnc: Function = _art_from_dict(func_types) - if fnc.name is None: - log_info(f"Function {item['function_id']} has no name, skipping signature application.") - continue - - addr = id_to_addr.get(result['origin_function_id']) - if not addr: - continue - log_info(f"RevEng.AI | Applying signature for {fnc.name} at 0x{addr:x}") - signature_data = {"deps": func_deps, "function": fnc} - apply_data_types_util(addr, signature_data, deci) - log_info(f"RevEng.AI | Successfully applied signature for {fnc.name} at 0x{addr:x}") - break - break - - except Exception as e: - log_error(f"RevEng.AI | Error resolving data types: {str(e)}") - - - def rename_functions(self, bv: BinaryView, options: List[Dict]): - try: - renamed_count = 0 - for option in options: - if rename_function_util( - self.config, - bv, - option['virtual_address'], - option['suggested_name'], - option['suggested_mangled_name'], - option['source_function_id'], - ): - renamed_count += 1 - return True, f"Successfully renamed {renamed_count} functions" - except Exception as e: - log_error(f"RevEng.AI | Error renaming functions: {str(e)}") - return False, str(e) - - def auto_unstrip(self, bv: BinaryView): - try: - log_info("RevEng.AI | Auto Unstripping binary") - - analysis_id = self.config.get_analysis_id(bv) - auto_unstrip_request = revengai.AutoUnstripRequest() - matches = [] - results = [] - - if self.cancelled.is_set(): - return False, "Operation cancelled" - - with self.config.create_api_client() as api_client: - api_instance = revengai.FunctionsCoreApi(api_client) - api_response = api_instance.auto_unstrip(analysis_id, auto_unstrip_request) - - if api_response.status.lower() == "error": - raise Exception(api_response.error_message) - elif api_response.status.lower() != "completed": - while True: - time.sleep(3) - api_response = api_instance.auto_unstrip(analysis_id, auto_unstrip_request) - if api_response.status.lower() == "completed": - break - if api_response.status.lower() == "error": - raise Exception(api_response.error_message) - - if api_response.status.lower() == "completed": - matches = api_response.matches - else: - raise Exception(api_response.error_message) - for match in matches: - print(match, flush=True) - try: - if self.cancelled.is_set(): - return False, "Operation cancelled" - - function = get_function_by_addr_util(bv, match.function_vaddr) - results.append({ - "virtual_address": match.function_vaddr, - "current_name": function.name, - "suggested_name": match.suggested_demangled_name, - "suggested_mangled_name": match.suggested_name, - "source_function_id": match.function_id, - }) - except Exception as e: - log_error(f"RevEng.AI | Error getting function by address {match.function_vaddr}: {str(e)}") - results.append({ - "virtual_address": match.function_vaddr, - "current_name": "N/A", - "suggested_name": match.suggested_demangled_name, - "suggested_mangled_name": match.suggested_name, - "source_function_id": match.function_id, - }) - - if self.cancelled.is_set(): - return False, "Operation cancelled" - - return True, results - - except Exception as e: - log_error(f"RevEng.AI | Error: {str(e)}") - return False, str(e) diff --git a/reai_toolkit/features/auto_unstrip/auto_unstrip_dialog.py b/reai_toolkit/features/auto_unstrip/auto_unstrip_dialog.py deleted file mode 100755 index 13058e2..0000000 --- a/reai_toolkit/features/auto_unstrip/auto_unstrip_dialog.py +++ /dev/null @@ -1,216 +0,0 @@ -import os -from binaryninja import log_error, log_info -from PySide6.QtGui import QPixmap -from PySide6.QtCore import Qt, QCoreApplication -from reai_toolkit.utils import create_progress_dialog, DataThread, create_cancellable_progress_dialog -from PySide6.QtWidgets import QDialog, QVBoxLayout, QHBoxLayout, QPushButton, QLabel, QCheckBox, QMessageBox, QTableWidget, QHeaderView, QAbstractItemView, QTableWidgetItem - -class AutoUnstripDialog(QDialog): - def __init__(self, config, auto_unstrip, bv): - super().__init__() - self.config = config - self.auto_unstrip = auto_unstrip - self.selected_results = [] - self.bv = bv - self.init_ui() - - - def init_ui(self): - self.setWindowTitle("RevEng.AI: Auto Unstrip Binary") - self.setMinimumSize(1000, 700) - self.resize(1200, 800) - - 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("Automatically rename unknown functions.") - description_label.setWordWrap(True) - info_layout.addWidget(title_label) - info_layout.addWidget(description_label) - header_layout.addLayout(info_layout, stretch=1) - - layout.addLayout(header_layout) - layout.addSpacing(20) - - self.results_table = QTableWidget() - self.results_table.setColumnCount(4) - self.results_table.setHorizontalHeaderLabels([ - "Select", "Virtual Address", "Current Name", "Suggested Name", - ]) - self.results_table.setSelectionMode(QAbstractItemView.NoSelection) - - header = self.results_table.horizontalHeader() - header.setSectionResizeMode(0, QHeaderView.ResizeToContents) - header.setSectionResizeMode(1, QHeaderView.ResizeToContents) - header.setSectionResizeMode(2, QHeaderView.Stretch) - header.setSectionResizeMode(3, QHeaderView.Stretch) - self.results_table.setAlternatingRowColors(True) - self.results_table.verticalHeader().setVisible(False) - layout.addWidget(self.results_table) - - layout.addSpacing(20) - - button_layout = QHBoxLayout() - self.save_button = QPushButton("Rename Functions") - self.save_button.setStyleSheet(""" - QPushButton { - background-color: #007bff; - color: white; - padding: 8px 16px; - border-radius: 4px; - } - QPushButton:hover { - background-color: #4400ff; - } - """) - self.save_button.clicked.connect(self._rename) - self.cancel_button = QPushButton("Cancel") - self.cancel_button.setStyleSheet(""" - QPushButton { - padding: 8px 16px; - border-radius: 4px; - } - """) - self.cancel_button.clicked.connect(self.reject) - - button_layout.addWidget(self.save_button) - button_layout.addWidget(self.cancel_button) - layout.addLayout(button_layout) - self.setLayout(layout) - self.show() - QCoreApplication.processEvents() - self._auto_unstrip() - - - def _rename(self): - self.progress = create_progress_dialog(self, "RevEng.AI Auto Unstrip", "Renaming selected functions...") - self.progress.show() - QCoreApplication.processEvents() - - self.rename_thread = DataThread(self.auto_unstrip.rename_functions, self.bv, self.selected_results) - self.rename_thread.finished.connect(self._on_rename_finished) - self.rename_thread.start() - - - def _on_rename_finished(self, success, message): - self.progress.close() - - if success: - QMessageBox.information(self, "RevEng.AI Auto Unstrip", message, QMessageBox.Ok) - self.accept() - else: - log_error(f"RevEng.AI | Failed to rename functions: {message}") - QMessageBox.critical(self, "RevEng.AI Auto Unstrip", f"Failed to rename functions: {message}", QMessageBox.Ok) - self.reject() - - - - def _auto_unstrip(self): - self.progress = create_cancellable_progress_dialog(self, "RevEng.AI Auto Unstrip", "Auto Unstripping binary...", self.auto_unstrip.cancel) - self.progress.show() - QCoreApplication.processEvents() - - self.auto_unstrip_thread = DataThread(self.auto_unstrip.auto_unstrip, self.bv, callback_cancelled_reset = self.auto_unstrip.clear_cancelled) - self.auto_unstrip_thread.finished.connect(self._on_auto_unstrip_finished) - self.auto_unstrip_thread.start() - - def _on_auto_unstrip_finished(self, success, data): - self.progress.close() - - if success: - self.populate_results_table(data) - message = f"Total functions found: {len(data)}" - else: - message = data - 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) - self.reject() - - def populate_results_table(self, results): - self.results_table.setRowCount(len(results)) - for i, result in enumerate(results): - select_item = QTableWidgetItem() - select_item.setFlags(Qt.ItemIsUserCheckable | Qt.ItemIsEnabled | Qt.ItemIsSelectable) - select_item.setCheckState(Qt.Checked) - self.results_table.setItem(i, 0, select_item) - - columns = [ - (1, "virtual_address", lambda x: f"0x{x:x}" if isinstance(x, int) else str(x)), - (2, "current_name", lambda x: str(x)), - (3, "suggested_name", lambda x: str(x)), - ] - - for col_idx, field, transform in columns: - value = result.get(field, "") - try: - safe_value = transform(value) - except Exception: - safe_value = str(value) - item = QTableWidgetItem(safe_value) - item.setFlags(item.flags() & ~Qt.ItemIsEditable) - item.setData(Qt.UserRole, result) - log_info(f"RevEng.AI | row: {i}, column: {col_idx}, item: {item.data(Qt.UserRole)}") - item.setSelected(False) - self.results_table.setItem(i, col_idx, item) - - - self.results_table.cellClicked.connect(self.on_checkbox_changed) - try: - self.results_table.cellClicked.disconnect() - except: - pass - self.results_table.cellClicked.connect(self.on_checkbox_changed) - - self.selected_results.append(result) - - 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 - - log_info(f"RevEng.AI | on_checkbox_changed called for row: {row}, column: {column}, is_checkbox: {is_checkbox}") - - result = self.results_table.item(row, 1).data(Qt.UserRole) - log_info(f"RevEng.AI | result: {result}") - result_virtual_address = result.get("virtual_address") if result else None - log_info(f"RevEng.AI | result_virtual_address: {result_virtual_address}") - - if result and result_virtual_address is not None: - log_info(f"RevEng.AI | result_virtual_address: {result_virtual_address}") - is_selected = any(r.get("virtual_address") == result_virtual_address for r in self.selected_results) - log_info(f"RevEng.AI | is_selected: {is_selected}") - - checkbox_item = self.results_table.item(row, 0) - current_state = checkbox_item.checkState() - - # Clear row selection - for col in range(self.results_table.columnCount()): - row_item = self.results_table.item(row, col) - if row_item: - row_item.setSelected(False) - - if current_state == Qt.Unchecked: - log_info(f"RevEng.AI | Removing result from selection: {result_virtual_address}") - self.selected_results = [r for r in self.selected_results if r.get("virtual_address") != result_virtual_address] - else: - log_info(f"RevEng.AI | Adding result to selection: {result_virtual_address}") - if not is_selected: - self.selected_results.append(result) - - - log_info(f"RevEng.AI | Total selected results: {len(self.selected_results)}") diff --git a/reai_toolkit/features/choose_source/choose_source_dialog.py b/reai_toolkit/features/choose_source/choose_source_dialog.py index 502d195..c461873 100755 --- a/reai_toolkit/features/choose_source/choose_source_dialog.py +++ b/reai_toolkit/features/choose_source/choose_source_dialog.py @@ -33,7 +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") + description_label = QLabel("Choose the source for your binary analysis. This selection will be used for all subsequent features in the plugin, including function analysis and other operations.\n\n") description_label.setWordWrap(True) info_layout.addWidget(title_label) info_layout.addWidget(description_label) diff --git a/reai_toolkit/features/configuration/__init__.py b/reai_toolkit/features/configuration/__init__.py index c116a46..8373c65 100755 --- a/reai_toolkit/features/configuration/__init__.py +++ b/reai_toolkit/features/configuration/__init__.py @@ -1,11 +1,13 @@ from reai_toolkit.features.configuration.config import Config from reai_toolkit.features.configuration.config_dialog import ConfigDialog +from reai_toolkit.utils import DataThread from PySide6.QtWidgets import QMessageBox from binaryninja import PluginCommand, log_info, BinaryViewType, log_error class ConfigurationFeature(): def __init__(self): self.config = Config() + self._init_threads = [] self._register_binary_view_event() log_info("RevEng.AI | Configuration Feature initialized") @@ -33,21 +35,27 @@ def _add_binaryview_finalized_event(self, bv): try: if bv.view_type == "Raw": return - + log_info(f"RevEng.AI | Binary view finalized: {bv.file.filename}") - status, message = self.config.init_config(bv) - if status: - log_info("RevEng.AI | Configuration initialized successfully") - elif message == "Binary not found in RevEng.AI, try processing the binary again.": - QMessageBox.warning( - 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 > Create new' option " - "before using other RevEng.AI features.", - QMessageBox.Ok - ) - else: - log_error(f"RevEng.AI | Configuration initialization failed: {message}") + thread = DataThread(self.config.init_config, bv) + thread.finished.connect(self._on_init_config_finished) + thread.finished.connect(lambda *_: self._init_threads.remove(thread)) + self._init_threads.append(thread) + thread.start() except Exception as e: log_error(f"RevEng.AI | Error in binary view event handler: {str(e)}") + + def _on_init_config_finished(self, status, message): + if status: + log_info("RevEng.AI | Configuration initialized successfully") + elif message == "Binary not found in RevEng.AI, try processing the binary again.": + QMessageBox.warning( + 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 > Create new' option " + "before using other RevEng.AI features.", + QMessageBox.Ok + ) + else: + log_error(f"RevEng.AI | Configuration initialization failed: {message}") diff --git a/reai_toolkit/features/configuration/config.py b/reai_toolkit/features/configuration/config.py index 76a908f..eb17dcf 100755 --- a/reai_toolkit/features/configuration/config.py +++ b/reai_toolkit/features/configuration/config.py @@ -15,6 +15,9 @@ def __init__(self): self.sha256 = "" self.api_config = None self.is_configured = False + self.binary_id = None + self.analysis_id = None + self.model_id = None self.user_agent = "Binary Ninja/RevEng.AI_Plugin" self.reveng_header = "X-RevEng-Application" self._load_config() diff --git a/reai_toolkit/revengai.py b/reai_toolkit/revengai.py index 75dfbdb..d8c4cbe 100755 --- a/reai_toolkit/revengai.py +++ b/reai_toolkit/revengai.py @@ -1,7 +1,6 @@ from binaryninja import log_info from reai_toolkit.features import ConfigurationFeature from reai_toolkit.features import UploadFeature -from reai_toolkit.features import AutoUnstripFeature from reai_toolkit.features import ChooseSourceFeature from reai_toolkit.features import MatchFunctionsFeature from reai_toolkit.features import MatchCurrentFunctionFeature @@ -14,7 +13,6 @@ def __init__(self): log_info("RevEng.AI | Initializing plugin") self.config_feature = ConfigurationFeature() 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()) @@ -28,7 +26,6 @@ def _register_features(self): log_info("RevEng.AI | Registering features") self.config_feature.register() 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() diff --git a/scripts/emit_features.py b/scripts/emit_features.py index 3624add..a67b4ca 100644 --- a/scripts/emit_features.py +++ b/scripts/emit_features.py @@ -7,7 +7,6 @@ FEATURE_CLASSES = { "UploadFeature": ("analyse_binary", "yes"), - "AutoUnstripFeature": ("binary_auto_analysis", "yes"), "ChooseSourceFeature": ("apply_existing_analysis", "yes"), "MatchCurrentFunctionFeature": ("rename_from_similar_function", "yes"), "AIDecompilerFeature": ("function_decompilation", "yes"), diff --git a/tests/headless/test_auto_unstrip.py b/tests/headless/test_auto_unstrip.py deleted file mode 100644 index f015a2c..0000000 --- a/tests/headless/test_auto_unstrip.py +++ /dev/null @@ -1,59 +0,0 @@ -from unittest.mock import MagicMock - -import pytest - -pytestmark = pytest.mark.headless -pytest.importorskip("binaryninja") - -from reai_toolkit.features.auto_unstrip import auto_unstrip as au_mod -from reai_toolkit.utils.core import binary_ninja as bn_mod - - -def test_auto_unstrip_maps_real_function_names(bv, mocker): - func = next(iter(bv.functions)) - config = MagicMock() - config.get_analysis_id.return_value = 1 - feature = au_mod.AutoUnstrip(config) - - match = MagicMock() - match.function_vaddr = func.start - match.suggested_demangled_name = "recovered" - match.suggested_name = "_Z9recoveredv" - match.function_id = 7 - api = mocker.patch.object(au_mod.revengai, "FunctionsCoreApi").return_value - api.auto_unstrip.return_value = MagicMock(status="completed", matches=[match]) - - ok, results = feature.auto_unstrip(bv) - - assert ok is True - assert results[0]["virtual_address"] == func.start - assert results[0]["current_name"] == func.name - assert results[0]["suggested_name"] == "recovered" - assert results[0]["source_function_id"] == 7 - - -def test_auto_unstrip_rename_applies_symbol_to_real_function(bv, mocker): - target = next( - (f for f in bv.functions if not (f.symbol and f.symbol.auto is False)), None - ) - if target is None: - pytest.skip("no auto-named function available to rename") - addr = target.start - portal = mocker.patch.object(bn_mod, "_rename_in_portal") - feature = au_mod.AutoUnstrip(MagicMock()) - - ok, _ = feature.rename_functions( - bv, - [ - { - "virtual_address": addr, - "suggested_name": "demangled_x", - "suggested_mangled_name": "mangled_x", - "source_function_id": 3, - } - ], - ) - - assert ok is True - assert bv.get_function_at(addr).name == "mangled_x" - portal.assert_called_once() diff --git a/tests/unit/auto_unstrip/test_service.py b/tests/unit/auto_unstrip/test_service.py deleted file mode 100644 index e40b4cf..0000000 --- a/tests/unit/auto_unstrip/test_service.py +++ /dev/null @@ -1,82 +0,0 @@ -from unittest.mock import MagicMock - -import pytest - -pytest.importorskip("binaryninja") - -from reai_toolkit.features.auto_unstrip import auto_unstrip as au_mod - - -def _match(vaddr, suggested_demangled, suggested, function_id): - match = MagicMock() - match.function_vaddr = vaddr - match.suggested_demangled_name = suggested_demangled - match.suggested_name = suggested - match.function_id = function_id - return match - - -@pytest.fixture -def feature(): - config = MagicMock() - config.get_analysis_id.return_value = 7 - return au_mod.AutoUnstrip(config) - - -@pytest.fixture -def api(mocker): - inst = mocker.patch.object(au_mod.revengai, "FunctionsCoreApi").return_value - return inst - - -def test_auto_unstrip_maps_completed_matches(feature, api, mocker): - api.auto_unstrip.return_value = MagicMock( - status="completed", - matches=[_match(0x1000, "foo", "_Z3foov", 42)], - ) - func = MagicMock() - func.name = "sub_1000" - mocker.patch.object(au_mod, "get_function_by_addr_util", return_value=func) - - ok, results = feature.auto_unstrip(MagicMock()) - - assert ok is True - assert results == [ - { - "virtual_address": 0x1000, - "current_name": "sub_1000", - "suggested_name": "foo", - "suggested_mangled_name": "_Z3foov", - "source_function_id": 42, - } - ] - - -def test_auto_unstrip_handles_missing_local_function(feature, api, mocker): - api.auto_unstrip.return_value = MagicMock( - status="completed", - matches=[_match(0x2000, "bar", "_Z3barv", 99)], - ) - mocker.patch.object( - au_mod, "get_function_by_addr_util", side_effect=Exception("not found") - ) - - ok, results = feature.auto_unstrip(MagicMock()) - - assert ok is True - assert results[0]["current_name"] == "N/A" - assert results[0]["suggested_name"] == "bar" - - -def test_auto_unstrip_polls_until_completed(feature, api, mocker): - api.auto_unstrip.side_effect = [ - MagicMock(status="processing"), - MagicMock(status="completed", matches=[]), - ] - mocker.patch.object(au_mod.time, "sleep") - - ok, results = feature.auto_unstrip(MagicMock()) - - assert ok is True - assert results == [] - assert api.auto_unstrip.call_count == 2 diff --git a/tests/unit/sdk/test_sdk_schemas.py b/tests/unit/sdk/test_sdk_schemas.py index b4f2de1..d9a91e5 100644 --- a/tests/unit/sdk/test_sdk_schemas.py +++ b/tests/unit/sdk/test_sdk_schemas.py @@ -7,7 +7,6 @@ from revengai.models.app_api_rest_v2_functions_types_function import ( AppApiRestV2FunctionsTypesFunction, ) -from revengai.models.auto_unstrip_response import AutoUnstripResponse from revengai.models.basic import Basic from revengai.models.binary_search_result import BinarySearchResult from revengai.models.comments_data import CommentsData @@ -44,7 +43,6 @@ "start_functions_matching", "get_functions_matching_status", "get_functions_matches", - "auto_unstrip", ], "FunctionsRenamingHistoryApi": ["rename_function_id"], "FunctionsDataTypesApi": [ @@ -148,11 +146,6 @@ def test_analysis_matching_status_and_results_accept_analysis_id(): assert "analysis_id" in _params(getattr(revengai.AnalysesCoreApi, method)) -def test_auto_unstrip_accepts_request_kwargs(): - params = _params(revengai.FunctionsCoreApi.auto_unstrip) - assert {"analysis_id", "auto_unstrip_request"} <= params - - def test_rename_function_id_accepts_kwargs(): params = _params(revengai.FunctionsRenamingHistoryApi.rename_function_id) assert {"function_id", "function_rename"} <= params @@ -269,11 +262,6 @@ def test_matched_function_suggestion_has_plugin_fields(): } <= set(MatchedFunctionSuggestion.model_fields) -def test_auto_unstrip_response_has_status_and_matches(): - fields = set(AutoUnstripResponse.model_fields) - assert {"status", "matches", "error_message"} <= fields - - def test_workflow_progress_exposes_status(): assert "status" in WorkflowProgress.model_fields diff --git a/uv.lock b/uv.lock index 22308b7..72b67dc 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.12'", @@ -144,7 +144,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -345,7 +345,7 @@ requires-dist = [ { name = "pycparser", specifier = ">=2.22,<3" }, { name = "pydantic", specifier = ">=2.12.5" }, { name = "requests", specifier = ">=2.32" }, - { name = "revengai", specifier = ">=3.96.3" }, + { name = "revengai", specifier = ">=3.100.0" }, { name = "urllib3", specifier = ">=2.0.0,<2.3.0" }, ] @@ -674,7 +674,7 @@ wheels = [ [[package]] name = "revengai" -version = "3.96.3" +version = "3.100.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "lazy-imports" }, @@ -683,9 +683,9 @@ dependencies = [ { name = "typing-extensions" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/62/60/7276cdc650f22cd24b0da044123f99162d16669e22d4430251ef432d7311/revengai-3.96.3.tar.gz", hash = "sha256:5339577a9eeada8a65d1b48ab563c309e15f4fd959616a23fa52b4fb1d3572c5", size = 358270, upload-time = "2026-06-24T12:56:23.677Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d5/e7/a071390b2ddfbffbf10567e92f748f68e5dcec2f1aef9d7d24a18506bddc/revengai-3.100.0.tar.gz", hash = "sha256:947041949f4fd0feb1bb61cdc4b344ffbe3ef535bb3f0d9637f98c88b43c6f9a", size = 361553, upload-time = "2026-06-29T10:46:52.563Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e2/d0/21c41a015a9656e143206d5ceb4acda59177168b4d8ccc19617c2b35ad3f/revengai-3.96.3-py3-none-any.whl", hash = "sha256:c855cd6351f5ccc1c42060a54fe0c256dafd47d7a85c06c0e54c72af79f43628", size = 933562, upload-time = "2026-06-24T12:56:22.099Z" }, + { url = "https://files.pythonhosted.org/packages/76/63/e0f17b77c3924e813a1331e881fe0d2e8aef97a3b4f80f9ab1c346301906/revengai-3.100.0-py3-none-any.whl", hash = "sha256:454ada33cbc89605914c28c65295aa5cc92126dd40faf4236cb76fe942801743", size = 941110, upload-time = "2026-06-29T10:46:50.969Z" }, ] [[package]]