Merge pull request #4 from RevEngAI/dev_ai_decompiler

Feat: AI Decompiler

Author: anisiosts

requirements.txt

@@ -1,4 +1,4 @@
-reait>=1.0.1
+reait>=1.2.5
 requests>=2.25.1
 PySide6>=6.0.0 
 libbs==2.15.1

revengai/features/__init__.py

@@ -4,6 +4,16 @@
 from .choose_source import ChooseSourceFeature
 from .match_functions import MatchFunctionsFeature
 from .match_current_function import MatchCurrentFunctionFeature
-from .view_function_in_portal import ViewFunctionInPortalFeature
+from .view_function_in_portal import ViewFunctionInPortalFeature    
+from .ai_decompiler import AIDecompilerFeature
 
-__all__ = ['ConfigurationFeature', 'UploadFeature', 'AutoUnstripFeature', 'ChooseSourceFeature', 'MatchFunctionsFeature', 'MatchCurrentFunctionFeature', 'ViewFunctionInPortalFeature'] 
+__all__ = [
+    'ConfigurationFeature', 
+    'UploadFeature', 
+    'AutoUnstripFeature', 
+    'ChooseSourceFeature', 
+    'MatchFunctionsFeature', 
+    'MatchCurrentFunctionFeature', 
+    'ViewFunctionInPortalFeature', 
+    'AIDecompilerFeature'
+    ] 

revengai/features/ai_decompiler/__init__.py

@@ -0,0 +1,69 @@
+from PySide6.QtCore import Qt
+from binaryninjaui import UIContext
+from .ai_decompiler import AIDecompiler
+from PySide6.QtWidgets import QDockWidget
+from revengai.utils import BaseAuthFeature
+from .ai_decompiler_dialog import AIDecompilerDialog
+from binaryninja import PluginCommand, log_info, BinaryView, log_error
+
+class AIDecompilerFeature(BaseAuthFeature):
+    def __init__(self, config=None):
+        super().__init__(config)
+        self.ai_decompiler = AIDecompiler(config)
+        self.dock_widget = None
+        self.widget = None
+        log_info("RevEng.AI | AIDecompiler Feature initialized")
+
+    def register(self):
+        PluginCommand.register_for_address(
+            "RevEng.AI\\7 - AI Decompiler",
+            "Get the AI decompiler for the current function",
+            self.show_ai_decompiler_dialog,
+            self.is_valid
+        )
+        log_info("RevEng.AI | AIDecompiler Feature registered")
+
+    def show_ai_decompiler_dialog(self, bv: BinaryView, func):
+        try:
+            log_info("RevEng.AI | Opening AI Decompiler Dock")
+            
+            ctx = UIContext.activeContext()
+            if not ctx:
+                log_error("RevEng.AI | No active UI context.")
+                return
+
+            main_win = ctx.mainWindow()
+            if not main_win:
+                log_error("RevEng.AI | No main window found.")
+                return
+
+            if self.dock_widget is not None and self.dock_widget.parent() is not None:
+                self.dock_widget.raise_()
+                log_info(f"RevEng.AI | AI Decompiler Dock already open, adding tab 0x{func:x}")
+                if self.widget is not None:
+                    self.widget.pre_tab_setup(bv, func)
+                return
+            
+            self.dock_widget = QDockWidget("RevEng.AI | AI Decompiler", main_win)
+            self.widget = AIDecompilerDialog(self.config, self.ai_decompiler, bv, func)
+            self.dock_widget.setObjectName("RevEng.AI | AI Decompiler")
+            self.dock_widget.setWidget(self.widget)
+            main_win.addDockWidget(Qt.RightDockWidgetArea, self.dock_widget)
+            self.dock_widget.raise_()
+            
+            self.dock_widget.visibilityChanged.connect(self.on_dock_closed)
+            log_info("RevEng.AI | AI Decompiler Dock displayed.")
+
+        except Exception as e:
+            log_error(f"RevEng.AI | Error opening AI Decompiler Dock: {e}")
+            return 
+
+    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

@@ -0,0 +1,134 @@
+from typing import Dict, Optional, Callable
+from binaryninja import BinaryView, log_info, log_error
+from reait.api import RE_poll_ai_decompilation, RE_begin_ai_decompilation
+from revengai.utils import get_function_id_by_addr as get_function_id_by_addr_util, AddressChangeMonitor, AIDecompilerChecker
+
+class AIDecompiler:
+    def __init__(self, config):
+        self.config = config
+        self._current_checker = None
+        self._track_timer = None
+        self.tracking_enabled = False
+        self._address_monitor = None
+        self._timer_monitor = None
+        self.dialog = None
+
+    def stop_ai_decompiler(self):
+        try:
+            if self._current_checker:
+                self._current_checker.stop()
+                self._current_checker = None
+                log_info("RevEng.AI | Stopped AI decompiler")
+        except Exception as e:
+            log_error(f"RevEng.AI | Error stopping AI decompiler: {str(e)}")
+
+    def stop_tracking(self):
+        try:
+            if self._track_timer:
+                self._track_timer.stop()
+                self._track_timer = None
+                log_info("RevEng.AI | Stopped active line tracking")
+        except Exception as e:
+            log_error(f"RevEng.AI | Error stopping active line tracking: {str(e)}")
+
+    def start_address_tracking(self, callback: Optional[Callable] = None, use_timer: bool = True):
+        try:
+            self.stop_address_tracking()
+            self._address_monitor = AddressChangeMonitor(self.address_change_callback)
+            log_info("RevEng.AI | Started notification-based address tracking")
+            
+        except Exception as e:
+            log_error(f"RevEng.AI | Error starting address tracking: {str(e)}")
+
+    def stop_address_tracking(self):
+        try:
+            if self._address_monitor:
+                self._address_monitor.unregister()
+                self._address_monitor = None
+                log_info("RevEng.AI | Stopped notification-based address tracking")
+                
+            if self._timer_monitor:
+                self._timer_monitor.stop()
+                self._timer_monitor = None
+                log_info("RevEng.AI | Stopped timer-based address tracking")
+                
+        except Exception as e:
+            log_error(f"RevEng.AI | Error stopping address tracking: {str(e)}")
+
+    def set_address_tracking_callback(self, callback: Callable):
+        try:
+            if self._address_monitor:
+                self._address_monitor.set_callback(callback)
+            if self._timer_monitor:
+                self._timer_monitor.set_callback(callback)
+            else:
+                self.start_address_tracking(callback)
+        except Exception as e:
+            log_error(f"RevEng.AI | Error setting address tracking callback: {str(e)}")
+    
+    def address_change_callback(self, context, view, addr, change_type):
+        if change_type == "address_changed" and addr is not None:
+            log_info(f"RevEng.AI | Address changed to 0x{addr:x} - could trigger AI decompilation here")
+            if self.dialog:
+                log_info(f"RevEng.AI | Pre-tab setup for address 0x{addr:x}")
+                bv = view.getCurrentViewInterface().getData() 
+                self.dialog.pre_tab_setup(bv, addr)
+
+    def start_ai_decompiler(self, bv: BinaryView, options: Dict) -> None:
+        try:
+            if not self.tracking_enabled:
+                self.start_address_tracking(self.address_change_callback)
+                self.tracking_enabled = True
+            
+            log_info("RevEng.AI | Starting function searching in portal")
+            editor = options.get("editor")
+            tab_name = options.get("tab_name")
+            function = options.get("function")
+            callback = options.get("callback")
+            binary_id = self.config.get_binary_id(bv)
+            function_id = get_function_id_by_addr_util(bv, function.start, binary_id)
+
+            res = RE_poll_ai_decompilation(
+                function_id,
+                summarise=True,
+            ).json()
+            
+            if not res.get("status", False):
+                callback(editor, "AI Decompilation failed.")
+                return
+            
+            poll_status = res.get("data").get("status", "uninitialised")
+            log_info(f"RevEng.AI | Polling AI decompilation: {poll_status}")
+
+            if poll_status == "uninitialised":
+                log_info(f"RevEng.AI | Starting AI Decompilation for function at 0x{function.start:x}")
+
+                try:
+                    res2 = RE_begin_ai_decompilation(
+                        function_id
+                    ).json()
+                except Exception as e:
+                    log_error(f"RevEng.AI | Error beginning AI decompilation: {str(e)}")
+                    callback(editor, "AI Decompilation failed.")
+                    return
+                
+                if not res2.get("status", False):
+                    callback(editor, "AI Decompilation failed.")
+                    return
+
+                log_info("RevEng.AI | AI Decompilation started")
+                periodic_checker = AIDecompilerChecker()
+                periodic_checker.start_ai_decompiler_checking(function_id, callback, editor, tab_name)
+                self._current_checker = periodic_checker
+
+            if poll_status == "success":
+                log_info(f"RevEng.AI | AI Decompilation for function at 0x{function.start:x} is completed")
+                callback(editor, res.get("data").get("decompilation"))
+
+            if poll_status == "error":
+                log_info(f"RevEng.AI | AI Decompilation for function at 0x{function.start:x} failed")
+                callback(editor, "AI Decompilation failed.")
+
+        except Exception as e:
+            log_error(f"RevEng.AI | Error in AI decompiler: {str(e)}")
+            return False, str(e)

revengai/features/ai_decompiler/ai_decompiler_dialog.py

@@ -0,0 +1,123 @@
+from PySide6.QtCore import QCoreApplication
+from binaryninja import log_error, log_info
+from revengai.utils import create_progress_dialog, get_function_by_addr, CHighlighter
+from PySide6.QtWidgets import QVBoxLayout, QCheckBox, QWidget, QTabWidget, QPlainTextEdit
+
+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.initial_setup_done = False
+        self.init_ui()
+
+    def init_ui(self):
+        self.setWindowTitle("RevEng.AI: AI Decompiler")
+        layout = QVBoxLayout()
+        self.setLayout(layout)
+        layout.addWidget(self.tabs)
+        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):
+        try:
+            progress_dialog = create_progress_dialog(self, "RevEng.AI", "Setting up AI Decompiler...")
+            progress_dialog.show()
+            QCoreApplication.processEvents()
+
+            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()
+
+            editor = QPlainTextEdit()
+            editor.setPlainText("Starting AI Decompiler...")
+            editor.setReadOnly(True)
+
+            editor.show()
+            QCoreApplication.processEvents()
+
+            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")
+        
+        try:
+            self.ai_decompiler.stop_ai_decompiler()
+        except Exception as e:
+            log_error(f"RevEng.AI | Error stopping AI decompiler on tab close: {str(e)}")
+        
+        self.tabs.removeTab(index)
+        if self.tabs.count() == 1: 
+            self.tabs.setTabsClosable(False)
+
+    def closeEvent(self, event):
+        try:
+            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)}")
+        
+        super().closeEvent(event)
+
+    def edit_editor(self, editor, text):
+        editor.setPlainText(text)
+        CHighlighter(editor.document())
+
+
+

revengai/features/auto_unstrip/__init__.py

@@ -1,7 +1,7 @@
-from binaryninja import PluginCommand, log_info, BinaryView
 from .auto_unstrip import AutoUnstrip
-from .auto_unstrip_dialog import AutoUnstripDialog
 from revengai.utils import BaseAuthFeature
+from .auto_unstrip_dialog import AutoUnstripDialog
+from binaryninja import PluginCommand, log_info, BinaryView
 
 class AutoUnstripFeature(BaseAuthFeature):
     def __init__(self, config=None):