fix: lib, upload and minor fixes

Author: anisiosts

reai_toolkit/__init__.py

@@ -2,13 +2,28 @@
 import os
 import sys
 from binaryninja import log_info, log_error
+import importlib
+
+def delete_module(module_name):
+    to_delete = []
+    for module in sys.modules:
+        if module.startswith(module_name):
+            to_delete.append(module)
+    for module_to_delete in to_delete:
+        log_info(f"RevEng.AI | Deleting module: {module_to_delete}")
+        del sys.modules[module_to_delete]
+
+def import_module(module_name):
+    try:
+        importlib.import_module(module_name)
+        log_info(f"RevEng.AI | Imported module: {module_name}")
+    except Exception as e:
+        log_error(f"RevEng.AI | Error importing module: {e}")
 
-# Get the directory where this __init__.py file is located
 try:
     plugin_dir = os.path.dirname(os.path.abspath(__file__))
     log_info(f"RevEng.AI | Using __file__ for plugin directory")
 except NameError:
-    # If __file__ is not defined, try to get it from the module
     import reai_toolkit
     plugin_dir = os.path.dirname(os.path.abspath(reai_toolkit.__file__))
     log_info(f"RevEng.AI | Using module path for plugin directory")
@@ -19,26 +34,42 @@
 log_info(f"RevEng.AI | Vendor path: {vendor_path}")
 log_info(f"RevEng.AI | Vendor exists: {os.path.exists(vendor_path)}")
 
-# Check if vendor directory exists and list its contents
 if os.path.exists(vendor_path):
-    # Add vendor directory to the beginning of sys.path for priority
     if vendor_path not in sys.path:
         sys.path.insert(0, vendor_path)
         log_info(f"RevEng.AI | Added vendor directory to sys.path at position 0")
 
-    # Also use site.addsitedir to handle .pth files
     site.addsitedir(vendor_path)
 
-    # List some contents for verification
     try:
         contents = os.listdir(vendor_path)
         log_info(f"RevEng.AI | Vendor directory contains {len(contents)} items")
         log_info(f"RevEng.AI | Sample contents: {', '.join(contents[:5])}")
+        
+        modules = ["urllib3", "certifi", "revengai"]
+
+        for module in modules:
+            delete_module(module)
+            import_module(module)
+
+        import certifi
+
+        os.environ["SSL_CERT_FILE"] = certifi.where()
+        os.environ["REQUESTS_CA_BUNDLE"] = certifi.where()
+
+        modules.remove("revengai")
+        sys.path.remove(vendor_path)
+
+        for module in modules:
+            delete_module(module)
+            import_module(module)
+
+        sys.path.insert(0, vendor_path)
+        
     except Exception as e:
         log_error(f"RevEng.AI | Error listing vendor contents: {e}")
 else:
     log_error(f"RevEng.AI | ERROR: Vendor directory not found at {vendor_path}")
-    # List what's actually in the plugin directory
     try:
         plugin_contents = os.listdir(plugin_dir)
         log_info(f"RevEng.AI | Plugin directory contains: {', '.join(plugin_contents)}")

reai_toolkit/features/auto_unstrip/auto_unstrip.py

@@ -8,6 +8,7 @@
 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):
@@ -16,6 +17,15 @@ def __init__(self, config):
         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:
@@ -89,6 +99,9 @@ def auto_unstrip(self, bv: BinaryView):
             matches = []
             results = []
 
+            if self.cancelled.is_set():
+                return False, "Operation cancelled"
+
             with revengai.ApiClient(self.config.api_config) as api_client:
                 api_instance = revengai.FunctionsCoreApi(api_client)
                 api_response = api_instance.auto_unstrip(analysis_id, auto_unstrip_request)
@@ -110,12 +123,18 @@ def auto_unstrip(self, bv: BinaryView):
                     raise Exception(api_response.error)
             for match in matches:
                 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})
                 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})
 
+            if self.cancelled.is_set():
+                return False, "Operation cancelled"
+
             return True, results
 
         except Exception as e:

reai_toolkit/features/auto_unstrip/auto_unstrip_dialog.py

@@ -2,7 +2,7 @@
 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
+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):
@@ -118,11 +118,11 @@ def _on_rename_finished(self, success, message):
 
 
     def _auto_unstrip(self):
-        self.progress = create_progress_dialog(self, "RevEng.AI Auto Unstrip", "Auto Unstripping binary...")
+        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)
+        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()