dev: getting click event in dev

anisiosts · anisiosts · commit 01062f8065c3 · 2025-08-01T20:21:24.000+01:00
diff --git a/revengai/features/ai_decompiler/__init__.py b/revengai/features/ai_decompiler/__init__.py
@@ -41,7 +41,7 @@ def show_ai_decompiler_dialog(self, bv: BinaryView, func):
                 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)
+                    self.widget.pre_tab_setup(bv, func)
                 return
             
             self.dock_widget = QDockWidget("RevEng.AI | AI Decompiler", main_win)
diff --git a/revengai/features/ai_decompiler/ai_decompiler.py b/revengai/features/ai_decompiler/ai_decompiler.py
@@ -1,12 +1,12 @@
 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 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
 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 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
@@ -18,44 +18,116 @@
     Struct,
     Typedef,
 )
+from revengai.utils.periodic_check import PeriodicChecker
+from PySide6.QtWidgets import QPlainTextEdit
+from binaryninja import BinaryView
+from binaryninjaui import UIContext
+from PySide6.QtCore import QTimer
 
 class AIDecompiler:
     def __init__(self, config):
         self.config = config
+        self._current_checker = None
+        self._track_timer = None  # Add timer instance variable
 
+    def stop_ai_decompiler(self):
+        """Stop the current AI decompiler checking"""
+        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):
+        """Stop the active line tracking"""
+        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 get_ai_decompiler(self, bv: BinaryView, options: Dict) -> None:
+
+    def start_ai_decompiler(self, bv: BinaryView, options: Dict) -> None:
         """Match functions from the binary against RevEng.AI database"""
         try:
+            ClickMonitor()
+            
             log_info("RevEng.AI | Starting function searching in portal")
-            function_addr = options.get("function", None)
+            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)
 
-            functions_containing = bv.get_functions_containing(function_addr)
+            res = RE_poll_ai_decompilation(
+                function_id,
+                summarise=True,
+            ).json()
             
-            if not functions_containing:
-                log_error(f"RevEng.AI | Function not found at 0x{function_addr:x}")
-                raise Exception("Function not found at address")
+            if not res.get("status", False):
+                callback(editor, "AI Decompilation failed.")
+                return
             
-            function = functions_containing[0]
-            log_info(f"RevEng.AI | Function: {function.name} at 0x{function.start:x} (Clicked address: 0x{function_addr:x})")
+            poll_status = res.get("data").get("status", "uninitialised")
+            log_info(f"RevEng.AI | Polling AI decompilation: {poll_status}")
 
-            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"]
+            if poll_status == "uninitialised":
+                log_info(f"RevEng.AI | Starting AI Decompilation for function at 0x{function.start:x}")
 
-            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)
+                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")
+
+                # Create PeriodicChecker instance (it's now a QObject)
+                periodic_checker = PeriodicChecker()
+                
+                # Start the AI decompiler checking with proper parameters
+                periodic_checker.start_ai_decompiler_checking(function_id, callback, editor, tab_name)
+                
+                # Store reference to prevent garbage collection
+                self._current_checker = periodic_checker
 
-            return True, "Function found in portal"
+            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"))
 
         except Exception as e:
-            log_error(f"RevEng.AI | Error in function matching: {str(e)}")
-            return False, str(e)
+            log_error(f"RevEng.AI | Error in AI decompiler: {str(e)}")
+            return False, str(e)
+        
+from binaryninja import PluginCommand, BinaryView
+from binaryninjaui import UIContext, UIContextNotification
+
+class ClickMonitor(UIContextNotification):
+    def __init__(self):
+        log_info("RevEng.AI | ClickMonitor initialized")
+        super().__init__()
+        log_info(f"RevEng.AI | class ClickMonitor")
+        for nome in dir(self):
+            attr = getattr(self, nome)
+            if callable(attr) and not nome.startswith("__") and not nome.startswith("On"):
+                log_info(f"RevEng.AI | {nome}")
+        UIContext.registerNotification(self)
+
+    def OnViewChanged(self, context, view):
+        log_info(f"RevEng.AI | View changed: {view}")
+
+    def OnAddressChanged(self, context, view, addr):
+        log_info(f"RevEng.AI | User navigated to address: {hex(addr)}")
+
diff --git a/revengai/features/ai_decompiler/ai_decompiler_dialog.py b/revengai/features/ai_decompiler/ai_decompiler_dialog.py
@@ -1,14 +1,17 @@
 from binaryninja import log_error, log_info
 from PySide6.QtWidgets import (QDockWidget, QVBoxLayout, QHBoxLayout, 
-                             QPushButton, QLabel, QCheckBox, QWidget, QTabWidget)
+                             QPushButton, QLabel, QCheckBox, QWidget, QTabWidget, QPlainTextEdit)
 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 .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
 
 class AIDecompilerDialog(QWidget):
@@ -19,41 +22,80 @@ def __init__(self, config, ai_decompiler, bv, func):
         self.bv = bv
         self.func = func
         self.tabs = QTabWidget()
-        self.init_ui(func)
+        self.number_of_clicks = 0
+        self.init_ui(bv, func)
 
-    def init_ui(self, func):
+    def init_ui(self, bv, func):
         self.setWindowTitle("RevEng.AI: AI Decompiler")
         layout = QVBoxLayout()
         self.setLayout(layout)
         layout.addWidget(self.tabs)
-        self.add_tab(func)
+        self.pre_tab_setup(bv, func)
         self.tabs.tabCloseRequested.connect(self.close_tab)
 
-    def add_tab(self, tab_name):
-        tab_name = str(f"0x{tab_name:x}")
+    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()):
-            log_info(f"RevEng.AI | Tab {self.tabs.tabText(i)} == {tab_name}")
             if self.tabs.tabText(i) == tab_name:
-                self.tabs.setCurrentIndex(i)
+                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()
-        layout.addWidget(QLabel(f"{tab_name} UI goes here"))
+
+        editor = QPlainTextEdit()
+        editor.setPlainText("Starting AI Decompiler...")
+        editor.setReadOnly(True)
+
+        #CHighlighter(editor.document())
+        editor.show()
+        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)
     
     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:
+            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):
+        """Handle dialog close event"""
+        try:
+            # Stop any running AI decompiler
+            if hasattr(self, '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())
+
+
 
diff --git a/revengai/features/ai_decompiler/c_highlighter.py b/revengai/features/ai_decompiler/c_highlighter.py
@@ -0,0 +1,80 @@
+from PySide6.QtGui import QSyntaxHighlighter, QTextCharFormat, QColor, QFont
+from PySide6.QtCore import QRegularExpression
+
+class CHighlighter(QSyntaxHighlighter):
+    def __init__(self, document = None):
+        super().__init__(document)
+
+        def fmt(color, bold=False):
+            f = QTextCharFormat()
+            f.setForeground(QColor(color))
+            if bold:
+                f.setFontWeight(QFont.Bold)
+            return f
+
+        self.highlighting_rules = []
+
+        # Keywords
+        keywords = [
+            "int", "char", "float", "double", "void", "if", "else", "while", "for", "return",
+            "switch", "case", "break", "continue", "struct", "typedef", "const", "unsigned",
+            "signed", "long", "short", "static", "volatile", "enum", "union", "do", "goto", "sizeof"
+        ]
+        keyword_format = fmt("#569CD6", True)
+        for word in keywords:
+            pattern = QRegularExpression(rf"\b{word}\b")
+            self.highlighting_rules.append((pattern, keyword_format))
+
+        # Types (optional C standard types)
+        type_format = fmt("#4EC9B0")
+        types = ["uint32_t", "int64_t", "uint8_t", "size_t"]
+        for t in types:
+            pattern = QRegularExpression(rf"\b{t}\b")
+            self.highlighting_rules.append((pattern, type_format))
+
+        # Strings
+        string_format = fmt("#CE9178")
+        self.highlighting_rules.append((QRegularExpression(r'"[^"\\]*(\\.[^"\\]*)*"'), string_format))
+
+        # Char literals
+        char_format = fmt("#CE9178")
+        self.highlighting_rules.append((QRegularExpression(r"'.'"), char_format))
+
+        # Comments
+        comment_format = fmt("#6A9955")
+        self.highlighting_rules.append((QRegularExpression(r"//.*"), comment_format))
+        self.comment_start = QRegularExpression(r"/\*")
+        self.comment_end = QRegularExpression(r"\*/")
+        self.comment_format = comment_format
+
+        # Numbers
+        number_format = fmt("#B5CEA8")
+        self.highlighting_rules.append((QRegularExpression(r"\b[0-9]+[uUlL]*\b"), number_format))
+
+    def highlightBlock(self, text):
+        for pattern, fmt in self.highlighting_rules:
+            iterator = pattern.globalMatch(text)
+            while iterator.hasNext():
+                match = iterator.next()
+                start, length = match.capturedStart(), match.capturedLength()
+                self.setFormat(start, length, fmt)
+
+        # Multiline comment highlighting
+        self.setCurrentBlockState(0)
+        start_index = 0
+        if self.previousBlockState() != 1:
+            match = self.comment_start.match(text)
+            start_index = match.capturedStart() if match.hasMatch() else -1
+
+        while start_index >= 0:
+            match_end = self.comment_end.match(text, start_index)
+            end_index = match_end.capturedEnd() if match_end.hasMatch() else -1
+
+            if end_index == -1:
+                self.setCurrentBlockState(1)
+                comment_length = len(text) - start_index
+            else:
+                comment_length = end_index - start_index
+
+            self.setFormat(start_index, comment_length, self.comment_format)
+            start_index = self.comment_start.match(text, start_index + comment_length).capturedStart()
diff --git a/revengai/utils/__init__.py b/revengai/utils/__init__.py
@@ -1,7 +1,8 @@
 from .periodic_check import PeriodicChecker
 from .base_auth_feature import BaseAuthFeature
 from .progress_dialog import create_progress_dialog, create_cancellable_progress_dialog
-from .utils import rename_function, parse_date
+from .utils import rename_function, parse_date, get_function_by_addr, get_function_id_by_addr
 from .datatypes import apply_data_types
 
-__all__ = ['PeriodicChecker', 'BaseAuthFeature', 'create_progress_dialog', 'create_cancellable_progress_dialog', 'rename_function', 'parse_date', 'apply_data_types']
+
+__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']
diff --git a/revengai/utils/periodic_check.py b/revengai/utils/periodic_check.py
diff --git a/revengai/utils/utils.py b/revengai/utils/utils.py
