Merge pull request #3 from RevEngAI/dev_view_functions_in_portal

feat: view function in portal added!

Author: anisiosts

revengai/features/__init__.py

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

revengai/features/match_current_function/match_current_function.py

@@ -123,60 +123,9 @@ def match_functions(self, bv: BinaryView, options: Dict[str, Any]) -> List[Dict]
 
         except Exception as e:
             log_error(f"RevEng.AI | Error in function matching: {str(e)}")
-            raise
-        
-        def parse_date(date_str: str) -> str:
-            dt = datetime.strptime(date_str, "%Y-%m-%dT%H:%M:%S.%f")
-            return dt.strftime("%Y-%m-%d %H:%M:%S")
-
-        def fetch_results(api_func, label: str) -> List[Dict[str, Any]]:
-            try:
-                log_info(f"RevEng.AI | Query: {query}")
-                response = api_func(query=query, page=1, page_size=1024).json()
-                results = response.get("data", {}).get("results", [])
-                log_info(f"Found {len(results)} {label.lower()}s")
-                return results
-            
-            except Exception as e:
-                log_error(f"RevEng.AI | Getting information failed. Reason: {str(e)}")
-                return []
-
-        def build_items(items_list: List[Dict[str, Any]], item_type: str) -> List[Tuple]:
-            items = []
-            for item in items_list:
-                name_key = "collection_name" if item_type == "Collection" else "binary_name"
-                date_key = "last_updated_at" if item_type == "Collection" else "created_at"
-                id_key = "collection_id" if item_type == "Collection" else "binary_id"
-                icon = "lock.png" if item_type == "Collection" and item["scope"] == "PRIVATE" else \
-                       "unlock.png" if item_type == "Collection" else "file.png"
-                
-                items.append({
-                    "name": item[name_key],
-                    "icon": icon,
-                    "type": item_type,
-                    "date": parse_date(item[date_key]),
-                    "model_name": item["model_name"],
-                    "owner": item["owned_by"],
-                    "id": item[id_key]
-                })
-            return items
-
-        try:
-
-            log_info(f"RevEng.AI | Searching for collections with '{query or 'N/A'}'")
-
-            collections_data = fetch_results(RE_collections_search, "collection")
-            binaries_data = fetch_results(RE_binaries_search, "binary")
-
-            table_items = build_items(collections_data, "Collection")
-            table_items += build_items(binaries_data, "Binary")
-
-            return table_items
-
-        except Exception as e:
-            log_error("Getting collections failed. Reason: %s", str(e))
             return False, str(e)
 
+
     def rename_function(self, bv: BinaryView, selected_result: Dict) -> List[Dict]:
         try:
             log_info(f"RevEng.AI | Starting function renaming for {len(selected_result)} functions")

revengai/features/view_function_in_portal/__init__.py

@@ -0,0 +1,27 @@
+from binaryninja import PluginCommand, log_info, BinaryView
+from .view_function_in_portal import ViewFunctionInPortal
+from .view_function_in_portal_dialog import ViewFunctionInPortalDialog
+from revengai.utils import BaseAuthFeature
+
+class ViewFunctionInPortalFeature(BaseAuthFeature):
+    def __init__(self, config=None):
+        super().__init__(config)
+        self.view_function_in_portal = ViewFunctionInPortal(config)
+        log_info("RevEng.AI | ViewFunctionInPortal Feature initialized")
+
+    def register(self):
+        PluginCommand.register_for_address(
+            "RevEng.AI\\7 - View Function in Portal",
+            "View the current function in the RevEng.AI portal",
+            self.show_match_current_function_dialog,
+            self.is_valid
+        )
+        log_info("RevEng.AI | ViewFunctionInPortal Feature registered")
+
+    def show_match_current_function_dialog(self, bv: BinaryView, func):
+        log_info("RevEng.AI | Opening MatchCurrentFunction dialog")
+        dialog = ViewFunctionInPortalDialog(self.config, self.view_function_in_portal, bv, func)
+        dialog.exec_() 
+
+    def is_valid(self, bv: BinaryView, func):
+        return self.config.is_configured == True 

revengai/features/view_function_in_portal/view_function_in_portal.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 ViewFunctionInPortal:
+    def __init__(self, config):
+        self.config = config
+
+
+    def view_function_in_portal(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/view_function_in_portal/view_function_in_portal_dialog.py

@@ -0,0 +1,70 @@
+from binaryninja import log_error
+from PySide6.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, 
+                             QPushButton, QLabel, QCheckBox)
+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 ViewFunctionInPortalDialog(QDialog):
+    def __init__(self, config, view_function_in_portal, bv, func):
+        super().__init__()
+        self.config = config
+        self.view_function_in_portal = view_function_in_portal
+        self.bv = bv
+        self.func = func
+        self.init_ui()
+
+    def init_ui(self):
+        self.setWindowTitle("RevEng.AI: View Function in Portal")
+        self.setWindowModality(Qt.WindowModal)
+        self.setMinimumSize(400, 100)
+        self.resize(400, 100)
+        
+        layout = QVBoxLayout()
+        
+        title_label = QLabel("Searching...")
+        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)
+        
+        self.setLayout(layout)
+        
+        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;
+            }
+        """)
+
+        options = {
+            "function": self.func
+        }
+
+        self.view_function_in_portal_thread = DataThread(self.view_function_in_portal.view_function_in_portal, self.bv, options)
+        self.view_function_in_portal_thread.finished.connect(self._on_view_function_in_portal_finished)
+        self.view_function_in_portal_thread.start()
+
+    def _on_view_function_in_portal_finished(self, success, message):
+        if success:
+            self.accept()
+        else:
+            log_error(f"RevEng.AI | Failed: {message}")
+            QMessageBox.critical(self, "RevEng.AI View Function in Portal Error", f"Failed to find function in portal: {message}", QMessageBox.Ok)
+            self.reject()

revengai/revengai.py

@@ -5,6 +5,7 @@
 from .features import ChooseSourceFeature
 from .features import MatchFunctionsFeature
 from .features import MatchCurrentFunctionFeature
+from .features import ViewFunctionInPortalFeature
 
 class RevengAIPlugin:
     def __init__(self):
@@ -15,6 +16,7 @@ def __init__(self):
         self.choose_source_feature = ChooseSourceFeature(self.config_feature.get_config())
         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._register_features()
 
     def _register_features(self):
@@ -25,4 +27,4 @@ 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()