Merge pull request #2 from RevEngAI/dev_match_functions

Feat: Match functions & Match Current Function

Author: anisiosts

README.md

@@ -151,4 +151,4 @@ This plugin is released under the GPL-2.0 license.
 
 ## Disclaimer
 
-Binary Ninja is a trademark of Vector 35. This project is not affiliated with or endorsed by Vector 35.
+Binary Ninja is a trademark of Vector 35. This project isn't affiliated with or endorsed by Vector 35.

requirements.txt

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

revengai/features/auto_unstrip/__init__.py

@@ -22,4 +22,4 @@ def show_auto_unstrip_dialog(self, bv: BinaryView):
         log_info("RevEng.AI | Opening AutoUnstrip dialog")
         dialog = AutoUnstripDialog(self.config, self.auto_unstrip, bv)
         dialog.exec_()
-
+    

revengai/features/auto_unstrip/auto_unstrip.py

@@ -1,9 +1,22 @@
 from binaryninja import BinaryView, log_info, log_error, Symbol, SymbolType
-from reait.api import RE_authentication, RE_search, RE_nearest_symbols_batch, RE_analyze_functions, RE_name_score
+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
 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 AutoUnstrip:
     def __init__(self, config):
@@ -13,12 +26,63 @@ def __init__(self, config):
         self.path = None
         self.max_workers = 4 
 
-    def _process_batch(self, function_ids: List[int], id_to_addr: Dict[int, int], bv: BinaryView) -> Tuple[int, List[str]]:
+    def resolve_data_types(self, to_datatypes: List[Dict], id_to_addr: Dict[int, int], deci: DecompilerInterface, chunk_index: int) -> None:
+        try:
+            function_ids = set([result['nearest_neighbor_id'] for result in to_datatypes])
+            log_info(f"RevEng.AI | Resolving data types for {len(function_ids)} functions")
+            RE_functions_data_types(function_ids=list(function_ids))
+            
+            items = []
+            while True:
+                response = RE_functions_data_types_poll(    
+                    function_ids=list(function_ids),
+                ).json()
+                data = response.get("data", {})
+                items = data.get("items", [])
+                pending_count = sum(1 for item in items if item.get("status") == "pending")
+                log_info(f"RevEng.AI | [Chunk {chunk_index}] {pending_count} items still pending...")
+                if not pending_count:
+                    break
+                time.sleep(3)
+
+            for item in items:
+                log_info(f"RevEng.AI | Item: {item['function_id']}")
+                if item['status'] != "completed":
+                    continue
+                for result in to_datatypes:
+                    if result['nearest_neighbor_id'] == item['function_id']:
+
+                        
+                        signature = "N/A"
+                        item2 = item.get("data_types", {})
+                        func_types = item2.get("func_types", None)
+                        func_deps = item2.get("func_deps", [])
+                        log_info(f"RevEng.AI | Func types: {func_types}")
+                        if func_types is not None:
+                            fnc: Function = _art_from_dict(func_types)
+                            if fnc.name is None:
+                                log_info(f"Function {item['function_id']} has no name, skipping signature application.")
+                                continue
+
+                            addr = id_to_addr.get(result['origin_function_id'])
+                            if not addr:
+                                continue
+                            log_info(f"RevEng.AI | Applying signature for {fnc.name} at 0x{addr:x}")
+                            signature_data = {"deps": func_deps, "function": fnc}
+                            apply_data_types_util(addr, signature_data, deci)
+                            log_info(f"RevEng.AI | Successfully applied signature for {fnc.name} at 0x{addr:x}")
+                            break
+                        break
+
+        except Exception as e:
+            log_error(f"RevEng.AI | Error resolving data types: {str(e)}")
+
+    def _process_batch(self, function_ids: List[int], id_to_addr: Dict[int, int], bv: BinaryView, debug_symbols: bool, data_types: bool, deci: DecompilerInterface = None, chunk_index: int = 0) -> Tuple[int, List[str]]:
         try:
             functions_by_distance = RE_nearest_symbols_batch(
                 function_ids=function_ids,
                 distance=self.auto_unstrip_distance,
-                debug_enabled=True,
+                debug_enabled=debug_symbols,
                 nns=1
             ).json()["function_matches"]
 
@@ -30,6 +94,7 @@ def _process_batch(self, function_ids: List[int], id_to_addr: Dict[int, int], bv
             #log_info(f"RevEng.AI | Functions by score: {functions_by_score}")
             renamed_count = 0
             errors = []
+            to_datatypes = []
             for result in functions_by_distance:
                 try:
                     func_id = result['origin_function_id']
@@ -47,15 +112,18 @@ def _process_batch(self, function_ids: List[int], id_to_addr: Dict[int, int], bv
 
                     for function in functions_by_score:
                         if function['function_id'] == func_id:
-                           if function['box_plot']["average"] < 0.9:
+                           if function['box_plot']["average"] < 90:
                                 log_info(f"RevEng.AI | Function {function['function_id']} has a score of {function['box_plot']['average']:.2f} for name {new_name_mangled}, skipping")
                                 break
                            else:
                                 log_info(f"RevEng.AI | Function {function['function_id']} has a score of {function['box_plot']['average']:.2f} for name {new_name_mangled}, renaming")
+                                to_datatypes.append(result)
                                 if rename_function_util(bv, func_addr, new_name_mangled):
                                     renamed_count += 1
                                 break
-                    
+                           
+                    if to_datatypes:
+                        self.resolve_data_types(to_datatypes, id_to_addr, deci, chunk_index)
 
                 except Exception as e:
                     log_error(f"RevEng.AI | Error processing function {result['origin_function_id']}: {str(e)}")
@@ -66,10 +134,13 @@ def _process_batch(self, function_ids: List[int], id_to_addr: Dict[int, int], bv
         except Exception as e:
             return 0, [str(e)]
 
-    def auto_unstrip(self, bv: BinaryView): 
+    def auto_unstrip(self, bv: BinaryView, options: Dict): 
         try:    
             log_info("RevEng.AI | Auto Unstripping binary")
 
+            debug_symbols = options.get("debug_symbols", True)
+            data_types = options.get("data_types", False)
+
             self.base_addr = bv.image_base
             self.path = bv.file.filename
             binary_id = self.config.get_binary_id(bv)
@@ -97,9 +168,13 @@ def auto_unstrip(self, bv: BinaryView):
 
             total_renamed = 0
             all_errors = []
+            deci = None
+            if data_types:
+                deci = BinjaInterface(bv)
+
             with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                 future_to_chunk = {
-                    executor.submit(self._process_batch, chunk, id_to_addr, bv): i 
+                    executor.submit(self._process_batch, chunk, id_to_addr, bv, debug_symbols, data_types, deci, i): i 
                     for i, chunk in enumerate(chunks)
                 }
 

revengai/features/auto_unstrip/auto_unstrip_dialog.py

@@ -1,6 +1,6 @@
 from binaryninja import log_error
 from PySide6.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, 
-                             QPushButton, QLabel)
+                             QPushButton, QLabel, QCheckBox)
 from PySide6.QtCore import Qt
 from PySide6.QtGui import QPixmap
 from PySide6.QtCore import QCoreApplication
@@ -18,7 +18,7 @@ def __init__(self, config, auto_unstrip, bv):
         self.init_ui()
 
     def init_ui(self):
-        self.setWindowTitle("Auto Unstrip Binary")
+        self.setWindowTitle("RevEng.AI: Auto Unstrip Binary")
         self.setMinimumWidth(500)
 
         layout = QVBoxLayout()
@@ -43,6 +43,14 @@ def init_ui(self):
         header_layout.addLayout(info_layout, stretch=1)
 
         layout.addLayout(header_layout)
+
+        self.debug_symbols_checkbox = QCheckBox("Limit Matches to Debug Symbols")
+        self.debug_symbols_checkbox.setChecked(True)
+        layout.addWidget(self.debug_symbols_checkbox)
+        self.data_types_checkbox = QCheckBox("Get functions Data Types (renaming may take longer)")
+        self.data_types_checkbox.setChecked(False)
+        layout.addWidget(self.data_types_checkbox)
+
         layout.addSpacing(20)
 
         button_layout = QHBoxLayout()
@@ -76,9 +84,14 @@ def init_ui(self):
     def _auto_unstrip(self):
         self.progress = create_progress_dialog(self, "RevEng.AI Auto Unstrip", "Auto Unstripping binary...")
         self.progress.show()
-        QCoreApplication.processEvents()   
+        QCoreApplication.processEvents()  
+
+        options = {
+            "debug_symbols": self.debug_symbols_checkbox.isChecked(),
+            "data_types": self.data_types_checkbox.isChecked()
+        }
 
-        self.auto_unstrip_thread = DataThread(self.auto_unstrip.auto_unstrip, self.bv)
+        self.auto_unstrip_thread = DataThread(self.auto_unstrip.auto_unstrip, self.bv, options)
         self.auto_unstrip_thread.finished.connect(self._on_auto_unstrip_finished)
         self.auto_unstrip_thread.start()
 

revengai/features/choose_source/choose_source_dialog.py

@@ -16,7 +16,7 @@ def __init__(self, config, choose_source, bv):
         self.init_ui()
 
     def init_ui(self):
-        self.setWindowTitle("Choose Source")
+        self.setWindowTitle("RevEng.AI: Choose Source")
         self.setMinimumWidth(500)
         layout = QVBoxLayout()
 

revengai/features/configuration/config_dialog.py

@@ -16,7 +16,7 @@ def __init__(self, config):
         self.init_ui()
 
     def init_ui(self):
-        self.setWindowTitle("RevEng.AI Configuration Wizard")
+        self.setWindowTitle("RevEng.AI: Configuration Wizard")
         self.setMinimumWidth(500)
 
         layout = QVBoxLayout()

revengai/features/match_current_function/__init__.py

@@ -11,7 +11,7 @@ def __init__(self, config=None):
 
     def register(self):
         PluginCommand.register_for_address(
-            "RevEng.AI\\Match Current Function",
+            "RevEng.AI\\6 - Match Current Function",
             "Search and match the current function against RevEng.AI database",
             self.show_match_current_function_dialog,
             self.is_valid