fixed autounstrip to work like portal

Author: anisiosts

features/auto_unstrip/auto_unstrip.py

@@ -1,5 +1,5 @@
 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
+from reait.api import RE_authentication, RE_search, RE_nearest_symbols_batch, RE_analyze_functions, RE_name_score
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import List, Dict, Tuple
 import math
@@ -36,22 +36,28 @@ def _rename_function(self, bv: BinaryView, addr: int, new_name: str, new_name_ma
     def _process_batch(self, function_ids: List[int], id_to_addr: Dict[int, int], bv: BinaryView) -> Tuple[int, List[str]]:
         """Process a batch of function IDs and return the number of renamed functions"""
         try:
-            ret = RE_nearest_symbols_batch(
+            functions_by_distance = RE_nearest_symbols_batch(
                 function_ids=function_ids,
                 distance=self.auto_unstrip_distance,
                 debug_enabled=True,
                 nns=1
             ).json()["function_matches"]
 
+            functions = []
+            for function in functions_by_distance:
+                functions.append({"function_id": function['origin_function_id'], "function_name": function['nearest_neighbor_function_name']})
+            log_info(f"RevEng.AI | Functions by distance: {functions}")
+            functions_by_score = RE_name_score(functions).json()["data"]
+            log_info(f"RevEng.AI | Functions by score: {functions_by_score}")
             renamed_count = 0
             errors = []
-            for result in ret:
+            for result in functions_by_distance:
                 try:
                     func_id = result['origin_function_id']
                     func_addr = id_to_addr.get(func_id)
                     if not func_addr:
                         continue
-                        
+
                     new_name = result['nearest_neighbor_function_name']
                     if not new_name or new_name.startswith(("sub_", "FUN_")):
                         continue
@@ -60,8 +66,18 @@ def _process_batch(self, function_ids: List[int], id_to_addr: Dict[int, int], bv
                     if not new_name_mangled or new_name_mangled.startswith(("sub_", "FUN_")):
                         continue
 
-                    if self._rename_function(bv, func_addr, new_name, new_name_mangled):
-                        renamed_count += 1
+                    for function in functions_by_score:
+                        if function['function_id'] == func_id:
+                           if function['box_plot']["average"] < 0.9:
+                                log_info(f"RevEng.AI | Function {function['function_id']} has a score of {function['box_plot']["average"]:.2f} for name {function['function_name']}, skipping")
+                                break
+                           else:
+                                log_info(f"RevEng.AI | Function {function['function_id']} has a score of {function['box_plot']["average"]:.2f} for name {function['function_name']}, renaming")
+                                if self._rename_function(bv, func_addr, new_name, new_name_mangled):
+                                    renamed_count += 1
+                                break
+                    
+
                 except Exception as e:
                     errors.append(str(e))
 
@@ -76,16 +92,17 @@ def auto_unstrip(self, bv: BinaryView):
 
             self.base_addr = bv.image_base
             self.path = bv.file.filename
+            binary_id = self.config.get_binary_id(bv)
             log_info(f"RevEng.AI | Path: {self.path}")
-            log_info(f"RevEng.AI | Binary ID: {self.config.binary_id}")
+            log_info(f"RevEng.AI | Binary ID: {binary_id}")
 
             results = RE_search(fpath=self.path).json()["query_results"]
             log_info(f"RevEng.AI | Search Results: {results}")
 
             if not len(results):
                 raise Exception("Binary not found in RevEng.AI, try uploading again.")
 
-            analyzed_functions = RE_analyze_functions(self.path, self.config.binary_id).json()["functions"]
+            analyzed_functions = RE_analyze_functions(self.path, binary_id).json()["functions"]
             function_ids = [func["function_id"] for func in analyzed_functions]
 
             id_to_addr = {

features/choose_source/choose_source.py

@@ -11,11 +11,14 @@ def __init__(self, config):
     def choose_source(self, bv: BinaryView, chose: str): 
         try:
             log_info(f"RevEng.AI | Chose: {chose}")
-            binary_id = chose.split("ID: ")[1].split(" -")[0]
-            log_info(f"RevEng.AI | Binary ID: {self.config.binary_id}")
-            # TODO: implement array like to store binary id and binaryhash and filesize
-
-            self.config.set_current_info(binary_id)
+            new_binary_id = int(chose.split("ID: ")[1].split(" -")[0])
+            binary_id = self.config.get_binary_id(bv)
+            if binary_id == new_binary_id:
+                log_info("RevEng.AI | Binary ID is already set to the chosen one.")
+                return True
+            
+            log_info(f"RevEng.AI | Changing Binary ID: {binary_id} to {new_binary_id}")
+            self.config.set_current_info(new_binary_id)
 
             return True
         except Exception as e:
@@ -25,9 +28,9 @@ def choose_source(self, bv: BinaryView, chose: str):
     def get_analysis(self, bv: BinaryView):
         try:
             log_info(f"RevEng.AI | Path: {bv.file.filename}")
-            log_info(f"RevEng.AI | Binary ID: {self.config.binary_id}")
-
-            # TODO: check if the binary is already in the database
+            binary_id = self.config.get_binary_id(bv)
+            log_info(f"RevEng.AI | Current Binary ID: {binary_id}")
+            
 
             results = RE_search(fpath=bv.file.filename).json()["query_results"]
 
@@ -37,9 +40,11 @@ def get_analysis(self, bv: BinaryView):
             options = []
             for result in results:
                 option = f"Name: {result['binary_name'][:10]}{'...' if len(result['binary_name']) > 10 else ''} - ID: {result['binary_id']} - Model: {result['model_name']} - Created at: {result['creation'].split('T')[0]} {result['creation'].split('T')[1].split('.')[0]}"
-                options.append(option)
                 log_info(f"RevEng.AI | Analysis: {option}")
-                # TODO: put the current binary id first in the list
+                if result['binary_id'] == binary_id:
+                    options.insert(0, option)
+                else:
+                    options.append(option)
 
             return options
         except Exception as e:

features/configuration/config.py

@@ -9,8 +9,6 @@ def __init__(self):
         self.host = ""
         self.sha256 = ""
         self.is_configured = False
-        self.binary_id = 0
-        self.analysis_id = 0
         self._load_config()
 
 
@@ -46,7 +44,7 @@ def save_config(self) -> bool:
             log_info(f"RevEng.AI | Testing configuration: {self.host} {self.api_key[:4]}...")
 
             if not self.check_auth():
-                raise(f"RevEng.AI | Authentication failed.")
+                raise Exception(f"RevEng.AI | Authentication failed.")
 
             self.is_configured = True
 
@@ -83,34 +81,34 @@ def get_all_analyses(self):
         all_analyses = settings.get_json("revengai.all_analyses", None)
         if all_analyses == "null":
             return {}
-        return all_analyses
+        return json.loads(all_analyses)
 
     def set_current_info(self, binary_id):
         binary_id = int(binary_id)
         self.binary_id = binary_id
 
         all_analyses = self.get_all_analyses()
         all_analyses[self.sha256] = {"binary_id": binary_id}
-        log_info(f"RevEng.AI | All analyses: {all_analyses}")
+        #log_info(f"RevEng.AI | All analyses: {all_analyses}")
         settings = Settings()
 
         settings.set_json("revengai.all_analyses", json.dumps(all_analyses))
         log_info(f"RevEng.AI | Updated analysis for SHA256: {self.sha256[:8]}... with binary_id: {binary_id}")
 
-
     def init_config(self, bv: BinaryView):
         try:
             if not self.check_auth():
                 self.is_configured = False
-                raise(f"RevEng.AI | Authentication failed.")
+                raise Exception("RevEng.AI | Authentication failed.")
 
             self.is_configured = True
 
             self.sha256 = re_binary_id(bv.file.filename)
             log_info(f"RevEng.AI | SHA256: {self.sha256}")
 
-            all_analyses = self.get_all_analyses()  
+            all_analyses = self.get_all_analyses()
             if self.sha256 in all_analyses:
+                log_info(f"RevEng.AI | Binary found in saved configurations, binary_id: {all_analyses[self.sha256]['binary_id']}")
                 self.binary_id = all_analyses[self.sha256]["binary_id"]
             else:
                 search_results = RE_search(fpath=bv.file.filename).json()["query_results"]
@@ -123,3 +121,10 @@ def init_config(self, bv: BinaryView):
             log_error(f"RevEng.AI | Failed to initialize configuration: {str(e)}")
             return False, str(e)
 
+    def get_binary_id(self, bv: BinaryView):
+        self.sha256 = re_binary_id(bv.file.filename)
+        all_analyses = self.get_all_analyses()
+        if self.sha256 in all_analyses:
+            return all_analyses[self.sha256]["binary_id"]
+        else:
+            return 0