dev: start of ai decomp

Author: anisiosts

revengai/features/__init__.py

@@ -4,6 +4,7 @@
 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,27 @@
+from binaryninja import PluginCommand, log_info, BinaryView
+from .ai_decompiler import AIDecompiler
+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)
+        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):
+        log_info("RevEng.AI | Opening MatchCurrentFunction dialog")
+        dialog = AIDecompilerDialog(self.config, self.ai_decompiler, bv, func)
+        dialog.exec_() 
+
+    def is_valid(self, bv: BinaryView, func):
+        return self.config.is_configured == True 

revengai/features/ai_decompiler/ai_decompiler.py

@@ -0,0 +1,61 @@
+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
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import List, Dict, Tuple
+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
+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,
+)
+
+class AIDecompiler:
+    def __init__(self, config):
+        self.config = config
+
+
+    def get_ai_decompiler(self, bv: BinaryView, options: Dict) -> None:
+        """Match functions from the binary against RevEng.AI database"""
+        try:
+            log_info("RevEng.AI | Starting function searching in portal")
+            function_addr = options.get("function", None)
+
+            functions_containing = bv.get_functions_containing(function_addr)
+            
+            if not functions_containing:
+                log_error(f"RevEng.AI | Function not found at 0x{function_addr:x}")
+                raise Exception("Function not found at address")
+            
+            function = functions_containing[0]
+            log_info(f"RevEng.AI | Function: {function.name} at 0x{function.start:x} (Clicked address: 0x{function_addr:x})")
+
+            binary_id = self.config.get_binary_id(bv)
+            if not binary_id:
+                raise Exception("Analysis not found. Please choose one using 'Choose Source' feature.")
+            
+            analysis = RE_get_analysis_id_from_binary_id(binary_id).json()
+            analyzed_functions = RE_get_functions_from_analysis(analysis["analysis_id"]).json()["data"]["functions"]
+
+            analyzed_function = next((f for f in analyzed_functions if (f["function_vaddr"] + bv.image_base) == function.start), None)
+            if not analyzed_function:
+                log_error(f"RevEng.AI | Function {function.name} not found in analyzed functions")
+                raise Exception("Function not found in analyzed functions")
+            
+            url = f"https://portal.reveng.ai/function/{analyzed_function['function_id']}"
+            InteractionHandler().open_url(url)
+
+            return True, "Function found in portal"
+
+        except Exception as e:
+            log_error(f"RevEng.AI | Error in function matching: {str(e)}")
+            return False, str(e)

revengai/features/ai_decompiler/ai_decompiler_dialog.py

@@ -0,0 +1,54 @@
+from binaryninja import log_error
+from PySide6.QtWidgets import (QDockWidget, QVBoxLayout, QHBoxLayout, 
+                             QPushButton, QLabel, QCheckBox, QWidget)
+from PySide6.QtCore import Qt
+from PySide6.QtGui import QPixmap
+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):
+    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()
+
+    def init_ui(self):
+        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)
+
+        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;
+            }
+        """)

revengai/features/view_function_in_portal/__init__.py

@@ -11,7 +11,7 @@ def __init__(self, config=None):
 
     def register(self):
         PluginCommand.register_for_address(
-            "RevEng.AI\\7 - View Function in Portal",
+            "RevEng.AI\\8 - View Function in Portal",
             "View the current function in the RevEng.AI portal",
             self.show_match_current_function_dialog,
             self.is_valid

revengai/revengai.py

@@ -6,6 +6,7 @@
 from .features import MatchFunctionsFeature
 from .features import MatchCurrentFunctionFeature
 from .features import ViewFunctionInPortalFeature
+from .features import AIDecompilerFeature
 
 class RevengAIPlugin:
     def __init__(self):
@@ -17,6 +18,7 @@ def __init__(self):
         self.match_functions_feature = MatchFunctionsFeature(self.config_feature.get_config())
         self.match_current_function_feature = MatchCurrentFunctionFeature(self.config_feature.get_config())
         self.view_function_in_portal_feature = ViewFunctionInPortalFeature(self.config_feature.get_config())
+        self.ai_decompiler_feature = AIDecompilerFeature(self.config_feature.get_config())
         self._register_features()
 
     def _register_features(self):
@@ -27,4 +29,5 @@ def _register_features(self):
         self.choose_source_feature.register()
         self.match_functions_feature.register()
         self.match_current_function_feature.register()
-        self.view_function_in_portal_feature.register()
+        self.view_function_in_portal_feature.register()
+        self.ai_decompiler_feature.register()