dev: starting ai decompiler feature

Author: anisiosts

revengai/features/__init__.py

@@ -7,4 +7,13 @@
 from .view_function_in_portal import ViewFunctionInPortalFeature    
 from .ai_decompiler import AIDecompilerFeature
 
-__all__ = ['ConfigurationFeature', 'UploadFeature', 'AutoUnstripFeature', 'ChooseSourceFeature', 'MatchFunctionsFeature', 'MatchCurrentFunctionFeature', 'ViewFunctionInPortalFeature', 'AIDecompilerFeature'] 
+__all__ = [
+    'ConfigurationFeature', 
+    'UploadFeature', 
+    'AutoUnstripFeature', 
+    'ChooseSourceFeature', 
+    'MatchFunctionsFeature', 
+    'MatchCurrentFunctionFeature', 
+    'ViewFunctionInPortalFeature', 
+    'AIDecompilerFeature'
+    ] 

revengai/features/ai_decompiler/__init__.py

@@ -1,12 +1,17 @@
-from binaryninja import PluginCommand, log_info, BinaryView
+from binaryninja import PluginCommand, log_info, BinaryView, log_error
+from PySide6.QtWidgets import QWidget, QVBoxLayout, QTabWidget, QLabel, QDockWidget
 from .ai_decompiler import AIDecompiler
+from binaryninjaui import UIContext
+from PySide6.QtCore import Qt
 from .ai_decompiler_dialog import AIDecompilerDialog
 from revengai.utils import BaseAuthFeature
 
 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):
@@ -19,9 +24,43 @@ def register(self):
         log_info("RevEng.AI | AIDecompiler Feature registered")
 
     def show_ai_decompiler_dialog(self, bv: BinaryView, func):
-        log_info("RevEng.AI | Opening MatchCurrentFunction dialog")
-        dialog = AIDecompilerDialog(self.config, self.ai_decompiler, bv, func)
-        dialog.exec_() 
+        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.add_tab(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:
+            self.dock_widget = None
 
     def is_valid(self, bv: BinaryView, func):
         return self.config.is_configured == True 

revengai/features/ai_decompiler/ai_decompiler_dialog.py

@@ -1,54 +1,59 @@
-from binaryninja import log_error
+from binaryninja import log_error, log_info
 from PySide6.QtWidgets import (QDockWidget, QVBoxLayout, QHBoxLayout, 
-                             QPushButton, QLabel, QCheckBox, QWidget)
+                             QPushButton, QLabel, QCheckBox, QWidget, QTabWidget)
 from PySide6.QtCore import Qt
 from PySide6.QtGui import QPixmap
+from PySide6.QtWidgets import QWidget
 from PySide6.QtCore import QCoreApplication
 from PySide6.QtWidgets import QMessageBox
 from PySide6.QtWidgets import QProgressBar
 from revengai.utils import create_progress_dialog
 from revengai.utils.data_thread import DataThread
 import os
 
-class AIDecompilerDialog(QDockWidget):
+class AIDecompilerDialog(QWidget):
     def __init__(self, config, ai_decompiler, bv, func):
         super().__init__()
         self.config = config
         self.ai_decompiler = ai_decompiler
         self.bv = bv
         self.func = func
-        self.init_ui()
+        self.tabs = QTabWidget()
+        self.init_ui(func)
 
-    def init_ui(self):
+    def init_ui(self, func):
         self.setWindowTitle("RevEng.AI: AI Decompiler")
-        
-        # Create a widget to hold the layout
-        content_widget = QWidget()
-        layout = QVBoxLayout(content_widget)
-        
-        title_label = QLabel("Getting AI decompiler...")
-        title_label.setStyleSheet("font-size: 18px;")
-        layout.addWidget(title_label)
+        layout = QVBoxLayout()
+        self.setLayout(layout)
+        layout.addWidget(self.tabs)
+        self.add_tab(func)
+        self.tabs.tabCloseRequested.connect(self.close_tab)
+
+    def add_tab(self, tab_name):
+        tab_name = str(f"0x{tab_name:x}")
+        for i in range(self.tabs.count()):
+            log_info(f"RevEng.AI | Tab {self.tabs.tabText(i)} == {tab_name}")
+            if self.tabs.tabText(i) == tab_name:
+                self.tabs.setCurrentIndex(i)
+                return
+
+        log_info(f"RevEng.AI | Adding tab {tab_name}")
+        tab = QWidget()
+        layout = QVBoxLayout()
+        layout.addWidget(QLabel(f"{tab_name} UI goes here"))
+        tab.setLayout(layout)
+        index = self.tabs.addTab(tab, tab_name)
+        self.tabs.setCurrentIndex(index)
+
+        if self.tabs.count() > 1:
+            self.tabs.setTabsClosable(True)
+    
+    def close_tab(self, index):
+        log_info(f"RevEng.AI | Closing tab {index} of {self.tabs.count()} tabs")
+        self.tabs.removeTab(index)
+        if self.tabs.count() == 1: 
+            self.tabs.setTabsClosable(False)
 
-        progress_bar = QProgressBar()
-        progress_bar.setMinimumWidth(250)
-        progress_bar.setMinimumHeight(20)
-        layout.addWidget(progress_bar)
 
-        # Set the content widget
-        self.setWidget(content_widget)
 
-        self.setStyleSheet("""
-            QProgressBar {
-                border: 1px solid #cccccc;
-                border-radius: 4px;
-                text-align: center;
-                background-color: #f0f0f0;
-                min-width: 250px;
-                min-height: 20px;
-            }
-            QProgressBar::chunk {
-                background-color: #007bff;
-                border-radius: 3px;
-            }
-        """)
+