implementing choose source and loading functions

Author: anisiosts

features/__init__.py

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

features/autounstrip/__init__.py features/auto_unstrip/__init__.pyfeatures/autounstrip/__init__.py renamed to features/auto_unstrip/__init__.py

@@ -1,25 +1,25 @@
 from binaryninja import PluginCommand, log_info, BinaryView
-from .autounstrip import AutoUnstrip
-from .autounstrip_dialog import AutoUnstripDialog
+from .auto_unstrip import AutoUnstrip
+from .auto_unstrip_dialog import AutoUnstripDialog
 from revengai_bn.utils import BaseAuthFeature
 
 class AutoUnstripFeature(BaseAuthFeature):
     def __init__(self, config=None):
         super().__init__(config)
-        self.autounstrip = AutoUnstrip(config)
+        self.auto_unstrip = 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.show_auto_unstrip_dialog,
             self.is_valid
         )
         log_info("RevEng.AI | AutoUnstrip Feature registered")
 
-    def show_autounstrip_dialog(self, bv: BinaryView):
+    def show_auto_unstrip_dialog(self, bv: BinaryView):
         log_info("RevEng.AI | Opening AutoUnstrip dialog")
-        dialog = AutoUnstripDialog(self.config, self.autounstrip, bv)
+        dialog = AutoUnstripDialog(self.config, self.auto_unstrip, bv)
         dialog.exec_()
 

features/autounstrip/autounstrip.py features/auto_unstrip/auto_unstrip.pyfeatures/autounstrip/autounstrip.py renamed to features/auto_unstrip/auto_unstrip.py

@@ -10,19 +10,7 @@ def __init__(self, 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
+        self.max_workers = 4 
 
     def _rename_function(self, bv: BinaryView, addr: int, new_name: str, new_name_mangled: str) -> bool:
         try:
@@ -38,7 +26,7 @@ def _rename_function(self, bv: BinaryView, addr: int, new_name: str, new_name_ma
             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}")
+            log_info(f"RevEng.AI | Renamed function at {hex(addr)} to {new_name_mangled}")
             return True
 
         except Exception as e:
@@ -90,7 +78,6 @@ def auto_unstrip(self, bv: BinaryView):
             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}")
@@ -106,23 +93,19 @@ def auto_unstrip(self, bv: BinaryView):
                 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:
@@ -146,12 +129,4 @@ def auto_unstrip(self, bv: BinaryView):
 
         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}")
-            """
-
+            return False, str(e)

…atures/autounstrip/autounstrip_dialog.py …ures/auto_unstrip/auto_unstrip_dialog.pyfeatures/autounstrip/autounstrip_dialog.py renamed to features/auto_unstrip/auto_unstrip_dialog.py features/autounstrip/autounstrip_dialog.py renamed to features/auto_unstrip/auto_unstrip_dialog.py

@@ -8,14 +8,14 @@
 from PySide6.QtCore import QCoreApplication
 from PySide6.QtWidgets import QMessageBox
 from revengai_bn.utils import create_progress_dialog
-from .autounstrip_thread import AutoUnstripThread
+from .auto_unstrip_thread import AutoUnstripThread
 import os
 
 class AutoUnstripDialog(QDialog):
-    def __init__(self, config, autounstrip, bv):
+    def __init__(self, config, auto_unstrip, bv):
         super().__init__()
         self.config = config
-        self.autounstrip = autounstrip
+        self.auto_unstrip = auto_unstrip
         self.bv = bv
         self.init_ui()
 
@@ -52,7 +52,7 @@ def init_ui(self):
                 background-color: #4400ff;
             }
         """)
-        self.save_button.clicked.connect(self.auto_unstrip)
+        self.save_button.clicked.connect(self._auto_unstrip)
         self.cancel_button = QPushButton("Cancel")
         self.cancel_button.setStyleSheet("""
             QPushButton {
@@ -68,13 +68,13 @@ def init_ui(self):
 
         self.setLayout(layout)
 
-    def auto_unstrip(self):
+    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 = AutoUnstripThread(self.auto_unstrip, self.bv)
         self.auto_unstrip_thread.finished.connect(self._on_auto_unstrip_finished)
         self.auto_unstrip_thread.start()
 

…atures/autounstrip/autounstrip_thread.py …ures/auto_unstrip/auto_unstrip_thread.pyfeatures/autounstrip/autounstrip_thread.py renamed to features/auto_unstrip/auto_unstrip_thread.py features/autounstrip/autounstrip_thread.py renamed to features/auto_unstrip/auto_unstrip_thread.py

@@ -3,14 +3,14 @@
 class AutoUnstripThread(QThread):
     finished = Signal(bool, str)  # Signal for success/failure and error message
 
-    def __init__(self, autounstrip, bv):
+    def __init__(self, auto_unstrip, bv):
         super().__init__()
-        self.autounstrip = autounstrip
+        self.auto_unstrip = auto_unstrip
         self.bv = bv
 
     def run(self):
         try:
-            success, message = self.autounstrip.auto_unstrip(self.bv)
+            success, message = self.auto_unstrip.auto_unstrip(self.bv)
             if success:
                 self.finished.emit(True, message)
             else:

features/choose_source/__init__.py

@@ -0,0 +1,25 @@
+from binaryninja import PluginCommand, log_info, BinaryView
+from .choose_source import ChooseSource
+from .choose_source_dialog import ChooseSourceDialog
+from revengai_bn.utils import BaseAuthFeature
+
+class ChooseSourceFeature(BaseAuthFeature):
+    def __init__(self, config=None):
+        super().__init__(config)
+        self.choose_source = ChooseSource(config)
+        log_info("RevEng.AI | Choose Source Feature initialized")
+
+    def register(self):
+        PluginCommand.register(
+            "RevEng.AI\\Choose Source",
+            "Choose a source for the binary analysis",
+            self.show_choose_source_dialog,
+            self.is_valid
+        )
+        log_info("RevEng.AI | Choose Source Feature registered")
+
+    def show_choose_source_dialog(self, bv: BinaryView):
+        log_info("RevEng.AI | Opening Choose Source dialog")
+        dialog = ChooseSourceDialog(self.config, self.choose_source, bv)
+        dialog.exec_()
+

features/choose_source/analysis_load_thread.py

@@ -0,0 +1,21 @@
+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)) 

features/choose_source/choose_source.py

@@ -0,0 +1,45 @@
+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 ChooseSource:
+    def __init__(self, config):
+        self.config = config
+
+    def choose_source(self, bv: BinaryView, chose: str): 
+        try:
+            log_info(f"RevEng.AI | Chose: {chose}")
+            binary_id = chose.split("ID: ")[1].split(" -")[0]
+            self.config.binary_id = binary_id
+            log_info(f"RevEng.AI | Binary ID: {self.config.binary_id}")
+            # TODO: implement array like to store binary id and binaryhash and filesize
+            return True
+        except Exception as e:
+            log_error(f"RevEng.AI | Failed to choose source: {str(e)}")
+            return False
+
+    def get_analysis(self, bv: BinaryView):
+        try:
+            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}")
+
+            results = RE_search(fpath=self.path).json()["query_results"]
+
+            if not len(results):
+                raise Exception("Binary not found in RevEng.AI, try processing the binary again.")
+            
+            options = []
+            for result in results:
+                option = f"Name: {result['binary_name'][:10]}{'...' if len(result['binary_name']) > 10 else ''} - ID: {result['binary_id']} - Model: {result['model_name']} - Created at: {result['creation'].split('T')[0]} {result['creation'].split('T')[1].split('.')[0]}"
+                options.append(option)
+                log_info(f"RevEng.AI | Analysis: {option}")
+                # TODO: put the current binary id first in the list
+
+            return options
+        except Exception as e:
+            log_error(f"RevEng.AI | Failed to get analysis: {str(e)}")
+            return []