final draft of auto-unstrip

Author: anisiosts

features/__init__.py

@@ -1,4 +1,5 @@
 from .configuration import ConfigurationFeature
 from .upload import UploadFeature
+from .autounstrip import AutoUnstripFeature
 
-__all__ = ['ConfigurationFeature', 'UploadFeature'] 
+__all__ = ['ConfigurationFeature', 'UploadFeature', 'AutoUnstripFeature'] 

features/autounstrip/__init__.py

@@ -0,0 +1,25 @@
+from binaryninja import PluginCommand, log_info, BinaryView
+from .autounstrip import AutoUnstrip
+from .autounstrip_dialog import AutoUnstripDialog
+from revengai_bn.utils import BaseAuthFeature
+
+class AutoUnstripFeature(BaseAuthFeature):
+    def __init__(self, config=None):
+        super().__init__(config)
+        self.autounstrip = AutoUnstrip(config)
+        log_info("RevEng.AI | AutoUnstrip Feature initialized")
+
+    def register(self):
+        PluginCommand.register(
+            "RevEng.AI\\AutoUnstrip",
+            "Attempt to recover stripped function names",
+            self.show_autounstrip_dialog,
+            self.is_valid
+        )
+        log_info("RevEng.AI | AutoUnstrip Feature registered")
+
+    def show_autounstrip_dialog(self, bv: BinaryView):
+        log_info("RevEng.AI | Opening AutoUnstrip dialog")
+        dialog = AutoUnstripDialog(self.config, self.autounstrip, bv)
+        dialog.exec_()
+

features/autounstrip/autounstrip.py

@@ -0,0 +1,157 @@
+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
+
+class AutoUnstrip:
+    def __init__(self, config):
+        self.config = config
+        self.auto_unstrip_distance = 0.09999999999999998
+        self.base_addr = None
+        self.path = None
+        self._analysed_functions = {}
+        self._functions = []
+        self.function_ids = []
+        self.max_workers = 4  # Number of parallel threads
+
+    def _get_all_functions(self):
+        return list(self.bv.functions)
+
+    def _get_analysed_functions(self):
+        return self._analysed_functions
+
+    def _get_sync_analysed_ids_local(self):
+        return self.function_ids
+
+    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_mangled)
+            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:
+            ret = RE_nearest_symbols_batch(
+                function_ids=function_ids,
+                distance=self.auto_unstrip_distance,
+                debug_enabled=True,
+                nns=1
+            ).json()["function_matches"]
+
+            renamed_count = 0
+            errors = []
+            for result in ret:
+                try:
+                    func_id = result['origin_function_id']
+                    func_addr = id_to_addr.get(func_id)
+                    if not func_addr:
+                        continue
+                        
+                    new_name = result['nearest_neighbor_function_name']
+                    if not new_name or new_name.startswith(("sub_", "FUN_")):
+                        continue
+                    
+                    new_name_mangled = result['nearest_neighbor_function_name_mangled']
+                    if not new_name_mangled or new_name_mangled.startswith(("sub_", "FUN_")):
+                        continue
+                    
+                    if self._rename_function(bv, func_addr, new_name, new_name_mangled):
+                        renamed_count += 1
+                except Exception as e:
+                    errors.append(str(e))
+
+            return renamed_count, errors
+
+        except Exception as e:
+            return 0, [str(e)]
+
+    def auto_unstrip(self, bv: BinaryView): 
+        try:    
+            log_info("RevEng.AI | Auto Unstripping binary")
+
+            self.base_addr = bv.image_base
+            self.path = bv.file.filename
+            log_info(f"RevEng.AI | Path: {self.path}")
+            log_info(f"RevEng.AI | Binary ID: {self.config.binary_id}")
+            self._analysed_functions = {}
+
+            results = RE_search(fpath=self.path).json()["query_results"]
+            log_info(f"RevEng.AI | Search Results: {results}")
+
+            if not len(results):
+                raise Exception("Binary not found in RevEng.AI, try uploading again.")
+            
+            analyzed_functions = RE_analyze_functions(self.path, self.config.binary_id).json()["functions"]
+            function_ids = [func["function_id"] for func in analyzed_functions]
+
+            id_to_addr = {
+                func["function_id"]: func["function_vaddr"] + self.base_addr
+                for func in analyzed_functions
+            }
+
+            # Use fixed chunk size of 50
+            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}")
+
+            # Process chunks in parallel
+            total_renamed = 0
+            all_errors = []
+            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
+                # Submit all tasks
+                future_to_chunk = {
+                    executor.submit(self._process_batch, chunk, id_to_addr, bv): i 
+                    for i, chunk in enumerate(chunks)
+                }
+
+                # Wait for all tasks to complete and collect results
+                for future in as_completed(future_to_chunk):
+                    chunk_index = future_to_chunk[future]
+                    try:
+                        renamed_count, errors = future.result()
+                        total_renamed += renamed_count
+                        all_errors.extend(errors)
+                        log_info(f"RevEng.AI | Chunk {chunk_index} completed: renamed {renamed_count} functions")
+                    except Exception as e:
+                        log_error(f"RevEng.AI | Error processing chunk {chunk_index}: {str(e)}")
+
+            if total_renamed > 0:
+                message = f"Successfully renamed {total_renamed} functions!"
+            else:
+                message = "After analyzing the binary, no functions were found to be renamed."
+            
+            if all_errors:
+                message += f"\nEncountered {len(all_errors)} errors during processing."
+            
+            log_info(f"RevEng.AI | {message}")
+            return True, message
+
+        except Exception as e:
+            log_error(f"RevEng.AI | Error: {str(e)}")
+            return False, str(e)
+            """
+            self._functions = self._get_all_functions()
+            self._get_analysed_functions()
+            self.function_ids = self._get_sync_analysed_ids_local()
+
+            log_info(f"RevEng.AI | {bv.file.original_filename}")
+            """
+

features/autounstrip/autounstrip_dialog.py

@@ -0,0 +1,104 @@
+from binaryninja import BinaryView, PluginCommand, log_info, log_error
+from PySide6.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, 
+                             QPushButton, QLabel, QCheckBox, 
+                             QGroupBox, QRadioButton, QSpacerItem,
+                             QSizePolicy)
+from PySide6.QtCore import Qt
+from PySide6.QtGui import QPixmap, QIcon
+from PySide6.QtCore import QCoreApplication
+from PySide6.QtWidgets import QMessageBox
+from revengai_bn.utils import create_progress_dialog
+from .autounstrip_thread import AutoUnstripThread
+import os
+
+class AutoUnstripDialog(QDialog):
+    def __init__(self, config, autounstrip, bv):
+        super().__init__()
+        self.config = config
+        self.autounstrip = autounstrip
+        self.bv = bv
+        self.init_ui()
+
+    def init_ui(self):
+        self.setWindowTitle("Auto Unstrip Binary")
+        self.setMinimumWidth(500)
+
+        layout = QVBoxLayout()
+
+        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.setWordWrap(True)
+        info_layout.addWidget(title_label)
+        info_layout.addWidget(description_label)
+        
+        layout.addLayout(info_layout)
+        layout.addSpacing(20)
+
+        # Buttons
+        button_layout = QHBoxLayout()
+        self.save_button = QPushButton("Auto Unstrip")
+        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.auto_unstrip)
+        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)
+
+    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...")
+        
+        # Create and start upload thread
+        self.auto_unstrip_thread = AutoUnstripThread(self.autounstrip, 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
+            )
+            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
+            ) 

features/autounstrip/autounstrip_thread.py

@@ -0,0 +1,19 @@
+from PySide6.QtCore import QThread, Signal
+
+class AutoUnstripThread(QThread):
+    finished = Signal(bool, str)  # Signal for success/failure and error message
+    
+    def __init__(self, autounstrip, bv):
+        super().__init__()
+        self.autounstrip = autounstrip
+        self.bv = bv
+        
+    def run(self):
+        try:
+            success, message = self.autounstrip.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)) 

features/configuration/__init__.py

@@ -1,11 +1,11 @@
-from binaryninja import PluginCommand, log_info
+from binaryninja import PluginCommand, log_info, BinaryViewType, log_error
 from .config import Config
 from .config_dialog import ConfigDialog
-from revengai_bn.utils import BaseAuthFeature
 
 class ConfigurationFeature():
     def __init__(self):
         self.config = Config()
+        self._register_binary_view_event()
         log_info("RevEng.AI | Configuration Feature initialized")
 
     def register(self):
@@ -14,11 +14,23 @@ def register(self):
             "Configure RevEng.AI settings",
             self.show_configuration
         )
+        log_info("RevEng.AI | Configuration Feature registered")
 
     def show_configuration(self, bv):
         log_info("RevEng.AI | Opening configuration wizard")
         wizard = ConfigDialog(self.config)
         wizard.exec_()
 
     def get_config(self):
-        return self.config
+        return self.config  
+    
+    def _register_binary_view_event(self):
+        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):
+        try:
+            log_info(f"RevEng.AI | Binary view finalized: {bv.file.original_filename}")
+            self.config.init_config()
+        except Exception as e:
+            log_error(f"RevEng.AI | Error in binary view event handler: {str(e)}")