feat:ai decompiler

Author: anisiosts

revengai/features/ai_decompiler/__init__.py

@@ -1,5 +1,5 @@
 from binaryninja import PluginCommand, log_info, BinaryView, log_error
-from PySide6.QtWidgets import QWidget, QVBoxLayout, QTabWidget, QLabel, QDockWidget
+from PySide6.QtWidgets import QDockWidget
 from .ai_decompiler import AIDecompiler
 from binaryninjaui import UIContext
 from PySide6.QtCore import Qt
@@ -60,7 +60,10 @@ def show_ai_decompiler_dialog(self, bv: BinaryView, func):
 
     def on_dock_closed(self, visible):
         if not visible:
+            log_info("RevEng.AI | AI Decompiler Dock closed")
             self.dock_widget = None
+            self.widget = None
+            self.ai_decompiler.stop_address_tracking()
 
     def is_valid(self, bv: BinaryView, func):
         return self.config.is_configured == True 

revengai/features/ai_decompiler/ai_decompiler.py

@@ -1,78 +1,102 @@
 from binaryninja import log_error, log_info
-from PySide6.QtWidgets import (QDockWidget, QVBoxLayout, QHBoxLayout, 
-                             QPushButton, QLabel, QCheckBox, QWidget, QTabWidget, QPlainTextEdit)
-from PySide6.QtCore import Qt
-from PySide6.QtGui import QPixmap
-from PySide6.QtWidgets import QWidget
+from PySide6.QtWidgets import QVBoxLayout, QCheckBox, QWidget, QTabWidget, QPlainTextEdit
 from PySide6.QtCore import QCoreApplication
-from PySide6.QtWidgets import QMessageBox
-from PySide6.QtWidgets import QProgressBar
-from revengai.utils import create_progress_dialog
-from .c_highlighter import CHighlighter
-from revengai.utils.data_thread import DataThread
-from revengai.utils import get_function_by_addr
-from revengai.utils.periodic_check import PeriodicChecker
-import os
+from revengai.utils import create_progress_dialog, get_function_by_addr
+from revengai.utils.c_highlighter import CHighlighter
 
 class AIDecompilerDialog(QWidget):
     def __init__(self, config, ai_decompiler, bv, func):
         super().__init__()
         self.config = config
         self.ai_decompiler = ai_decompiler
+        self.ai_decompiler.dialog = self
         self.bv = bv
         self.func = func
         self.tabs = QTabWidget()
         self.number_of_clicks = 0
-        self.init_ui(bv, func)
+        self.initial_setup_done = False
+        self.init_ui()
 
-    def init_ui(self, bv, func):
+    def init_ui(self):
         self.setWindowTitle("RevEng.AI: AI Decompiler")
         layout = QVBoxLayout()
         self.setLayout(layout)
         layout.addWidget(self.tabs)
-        self.pre_tab_setup(bv, func)
+        self.address_tracking_checkbox = QCheckBox("Address Tracking")
+        self.address_tracking_checkbox.setChecked(True)
+        layout.addWidget(self.address_tracking_checkbox)
+        self.address_tracking_checkbox.stateChanged.connect(self.toggle_address_tracking)
         self.tabs.tabCloseRequested.connect(self.close_tab)
 
+    def showEvent(self, event):
+        super().showEvent(event)
+        if not self.initial_setup_done:
+            self.initial_setup_done = True
+            QCoreApplication.processEvents()
+            self.pre_tab_setup(self.bv, self.func)
+
+    def toggle_address_tracking(self, state):
+        log_info(f"RevEng.AI | Address tracking checkbox state changed to {state}")
+        if state == 2:
+            log_info(f"RevEng.AI | Starting address tracking")
+            self.ai_decompiler.start_address_tracking()
+        else:
+            log_info(f"RevEng.AI | Stopping address tracking")
+            self.ai_decompiler.stop_address_tracking()
+ 
     def pre_tab_setup(self, bv, func):
-        function = get_function_by_addr(bv, func)
-        tab_name = str(f"0x{function.start:x}")
-        log_info(f"RevEng.AI | Given address 0x{func:x} is function: {function.name} at 0x{function.start:x}")
-        for i in range(self.tabs.count()):
-            if self.tabs.tabText(i) == tab_name:
-                log_info(f"RevEng.AI | Tab {tab_name} already exists, switching to it.")
-                return
-            
-        log_info(f"RevEng.AI | Adding tab {tab_name}")
-        tab = QWidget()
-        layout = QVBoxLayout()
+        try:
+            progress_dialog = create_progress_dialog(self, "RevEng.AI", "Setting up AI Decompiler...")
+            progress_dialog.show()
+            QCoreApplication.processEvents()
 
-        editor = QPlainTextEdit()
-        editor.setPlainText("Starting AI Decompiler...")
-        editor.setReadOnly(True)
+            function = get_function_by_addr(bv, func)
+            tab_name = str(f"0x{function.start:x}")
+            log_info(f"RevEng.AI | Given address 0x{func:x} is function: {function.name} at 0x{function.start:x}")
+            for i in range(self.tabs.count()):
+                if self.tabs.tabText(i) == tab_name:
+                    log_info(f"RevEng.AI | Tab {tab_name} already exists, switching to it.")
+                    self.tabs.setCurrentIndex(i)
+                    progress_dialog.close()
+                    return
+                
+            log_info(f"RevEng.AI | Adding tab {tab_name}")
+            tab = QWidget()
+            layout = QVBoxLayout()
 
-        #CHighlighter(editor.document())
-        editor.show()
-        layout.addWidget(editor)
-        tab.setLayout(layout)
-        
-        index = self.tabs.addTab(tab, tab_name)
-        self.tabs.setCurrentIndex(index)
+            editor = QPlainTextEdit()
+            editor.setPlainText("Starting AI Decompiler...")
+            editor.setReadOnly(True)
 
-        if self.tabs.count() > 1:
-            self.tabs.setTabsClosable(True)
+            editor.show()
+            QCoreApplication.processEvents()
 
-        options = {
-            "editor": editor,
-            "tab_name": tab_name,
-            "function": function,
-            "callback": self.edit_editor
-        }
-        self.ai_decompiler.start_ai_decompiler(self.bv, options)
+            layout.addWidget(editor)
+            tab.setLayout(layout)
+            
+            index = self.tabs.addTab(tab, tab_name)
+            self.tabs.setCurrentIndex(index)
+
+            if self.tabs.count() > 1:
+                self.tabs.setTabsClosable(True)
+
+            options = {
+                "editor": editor,
+                "tab_name": tab_name,
+                "function": function,
+                "callback": self.edit_editor
+            }
+            
+            self.ai_decompiler.start_ai_decompiler(self.bv, options)
+            progress_dialog.close()
+        except Exception as e:
+            if 'progress_dialog' in locals():
+                progress_dialog.close()
+            log_error(f"RevEng.AI | Error during pre_tab_setup: {str(e)}")
 
     def close_tab(self, index):
         log_info(f"RevEng.AI | Closing tab {index} of {self.tabs.count()} tabs")
 
-        # Stop AI decompiler for this tab if it's running
         try:
             self.ai_decompiler.stop_ai_decompiler()
         except Exception as e:
@@ -83,10 +107,9 @@ def close_tab(self, index):
             self.tabs.setTabsClosable(False)
 
     def closeEvent(self, event):
-        """Handle dialog close event"""
         try:
-            # Stop any running AI decompiler
             if hasattr(self, 'ai_decompiler'):
+                log_info(f"RevEng.AI | Stopping AI decompiler")
                 self.ai_decompiler.stop_ai_decompiler()
         except Exception as e:
             log_error(f"RevEng.AI | Error during cleanup: {str(e)}")

revengai/features/ai_decompiler/ai_decompiler_dialog.py

@@ -26,28 +26,18 @@ def init_ui(self):
         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):
@@ -174,10 +164,8 @@ def start_matching(self):
         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,
@@ -187,13 +175,11 @@ def start_renaming(self):
             )
             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, 
@@ -205,8 +191,6 @@ def start_renaming(self):
     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_cancellable_progress_dialog(self, "RevEng.AI Fetch Data Types", "Fetching data types...", self.match_current_function.cancel)
             self.progress.show()
             QCoreApplication.processEvents() 
@@ -326,8 +310,3 @@ def on_fetching_data_types_finished(self, success, data):
                 f"Failed to fetch data types:\n{data}",
                 QMessageBox.Ok
             )
-
-    """
-    def closeEvent(self, event):
-        self.accept() 
-    """ 

revengai/features/match_current_function/match_current_function_dialog.py

@@ -3,6 +3,6 @@
 from .progress_dialog import create_progress_dialog, create_cancellable_progress_dialog
 from .utils import rename_function, parse_date, get_function_by_addr, get_function_id_by_addr
 from .datatypes import apply_data_types
+from .c_highlighter import CHighlighter
 
-
-__all__ = ['PeriodicChecker', 'BaseAuthFeature', 'create_progress_dialog', 'create_cancellable_progress_dialog', 'rename_function', 'parse_date', 'apply_data_types', 'get_function_by_addr', 'get_function_id_by_addr']
+__all__ = ['PeriodicChecker', 'BaseAuthFeature', 'create_progress_dialog', 'create_cancellable_progress_dialog', 'rename_function', 'parse_date', 'apply_data_types', 'get_function_by_addr', 'get_function_id_by_addr', 'CHighlighter']

revengai/utils/__init__.py

@@ -0,0 +1,98 @@
+from binaryninja import BinaryView, log_info, log_error, Symbol, SymbolType, interaction
+from binaryninja.interaction import InteractionHandler
+from reait.api import RE_authentication, RE_search, RE_nearest_symbols_batch, RE_analyze_functions, RE_name_score, RE_functions_data_types, RE_functions_data_types_poll, RE_get_analysis_id_from_binary_id, RE_get_functions_from_analysis, RE_poll_ai_decompilation, RE_begin_ai_decompilation
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import List, Dict, Tuple, Optional, Callable
+import math
+from revengai.utils.datatypes import apply_data_types as apply_data_types_util
+import time
+from revengai.utils import rename_function as rename_function_util, get_function_id_by_addr as get_function_id_by_addr_util
+from libbs.api import DecompilerInterface
+from libbs.decompilers.binja.interface import BinjaInterface
+from libbs.artifacts import _art_from_dict
+from libbs.artifacts import (
+    Function,
+    FunctionArgument,
+    GlobalVariable,
+    Enum,
+    Struct,
+    Typedef,
+)
+from revengai.utils.periodic_check import PeriodicChecker
+from PySide6.QtWidgets import QPlainTextEdit
+from binaryninja import BinaryView
+from binaryninjaui import UIContext, UIContextNotification
+from PySide6.QtCore import QTimer
+
+class AddressChangeMonitor(UIContextNotification):
+    """
+    Monitors address changes in Binary Ninja's UI and calls a callback function
+    when the user navigates to a new address in the decompiler view.
+    """
+    
+    def __init__(self, callback: Optional[Callable] = None):
+        """
+        Initialize the address change monitor.
+        
+        Args:
+            callback: Optional callback function to call when address changes.
+                     Should accept parameters: (context, view, address)
+        """
+        super().__init__()
+        self.callback = callback
+        self._registered = False
+        self._last_address = None
+        
+        # Register for notifications
+        self.register()
+        
+    def register(self):
+        """Register this notification with the UI context"""
+        if not self._registered:
+            UIContext.registerNotification(self)
+            self._registered = True
+            log_info("RevEng.AI | AddressChangeMonitor registered for notifications")
+    
+    def unregister(self):
+        """Unregister this notification from the UI context"""
+        if self._registered:
+            UIContext.unregisterNotification(self)
+            self._registered = False
+            log_info("RevEng.AI | AddressChangeMonitor unregistered from notifications")
+    
+    def set_callback(self, callback: Callable):
+        """Set or update the callback function"""
+        self.callback = callback
+        log_info("RevEng.AI | AddressChangeMonitor callback updated")
+
+    def OnAddressChange(self, context, view, frame, addr):
+        """Called when the user navigates to a new address"""
+        try:
+            # Extract the address from the ViewLocation object
+            current_addr = None
+            
+            # Handle ViewLocation object
+            if hasattr(addr, 'getOffset'):
+                current_addr = addr.getOffset()
+            elif hasattr(addr, 'addr'):
+                current_addr = addr.addr
+            elif hasattr(addr, 'address'):
+                current_addr = addr.address
+            else:
+                current_addr = addr
+                
+            
+            # Avoid duplicate notifications for the same address
+            if current_addr == self._last_address:
+                return
+                
+            self._last_address = current_addr
+            
+            if self.callback:
+                try:
+                    self.callback(context, view, current_addr, "address_changed")
+                except Exception as e:
+                    log_error(f"RevEng.AI | Error in address change callback: {str(e)}")
+        except Exception as e:
+            log_error(f"RevEng.AI | Error processing address change: {str(e)}")
+            log_info(f"RevEng.AI | OnAddressChange called with addr type: {type(addr)}")