fix: dynamic model id

Author: vgoat21

reai_toolkit/features/choose_source/choose_source.py

@@ -11,6 +11,7 @@ def choose_source(self, bv: BinaryView, chose: str):
             log_info(f"RevEng.AI | Item: {chose}")
             new_binary_id = int(chose.split("Binary ID: ")[1].split(" -")[0])
             new_analysis_id = int(chose.split("Analysis ID: ")[1].split(" -")[0])
+            new_model_id = int(chose.split("Model_ID: ")[1])
             analysis_id = self.config.get_analysis_id(bv)
             binary_id = self.config.get_binary_id(bv)
             if binary_id == new_binary_id and analysis_id == new_analysis_id:
@@ -19,7 +20,8 @@ def choose_source(self, bv: BinaryView, chose: str):
 
             log_info(f"RevEng.AI | Changing Binary ID: {binary_id} to {new_binary_id}")
             log_info(f"RevEng.AI | Changing Analysis ID: {analysis_id} to {new_analysis_id}")
-            self.config.set_current_info(new_binary_id, new_analysis_id)
+            log_info(f"RevEng.AI | Changing Model ID to {new_model_id}")
+            self.config.set_current_info(new_binary_id, new_analysis_id, new_model_id)
 
             return True, "Binary ID changed successfully."
         except Exception as e:
@@ -45,7 +47,7 @@ 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 ''} - Binary ID: {result.binary_id} - Analysis ID: {result.analysis_id} - Model: {result.model_name} - Created at: {result.created_at.strftime('%m/%d/%Y %H:%M')}"
+                option = f"Name: {result.binary_name[:10]}{'...' if len(result.binary_name) > 10 else ''} - Binary ID: {result.binary_id} - Analysis ID: {result.analysis_id} - Model: {result.model_name} - Created at: {result.created_at.strftime('%m/%d/%Y %H:%M')} - Model_ID: {result.model_id}"
                 log_info(f"RevEng.AI | Analysis: {option}")
                 if result.binary_id == binary_id and result.analysis_id == analysis_id:
                     options.insert(0, option)

reai_toolkit/features/configuration/config.py

@@ -105,16 +105,18 @@ def get_all_analyses(self):
         return json.loads(all_analyses)
 
 
-    def set_current_info(self, binary_id, analysis_id):
+    def set_current_info(self, binary_id, analysis_id, model_id):
 
         try:
             binary_id = int(binary_id)
             analysis_id = int(analysis_id)
+            model_id = int(model_id)
             self.binary_id = binary_id
             self.analysis_id = analysis_id
+            self.model_id = model_id
 
             all_analyses = self.get_all_analyses()
-            all_analyses[self.sha256] = {"binary_id": binary_id, "analysis_id": analysis_id}
+            all_analyses[self.sha256] = {"binary_id": binary_id, "analysis_id": analysis_id, "model_id": model_id}
             settings = Settings()
 
             settings.set_json("revengai.all_analyses", json.dumps(all_analyses))
@@ -142,6 +144,7 @@ def init_config(self, bv: BinaryView):
                 log_info(f"RevEng.AI | Binary found in saved configurations, binary_id: {all_analyses[self.sha256]['binary_id']} and analysis_id: {all_analyses[self.sha256]['analysis_id']}")
                 self.binary_id = all_analyses[self.sha256]["binary_id"]
                 self.analysis_id = all_analyses[self.sha256]["analysis_id"]
+                self.model_id = all_analyses[self.sha256]["model_id"]
             else:
                 log_info(f"RevEng.AI | Binary not found in saved configurations, searching in RevEng.AI...")
                 with revengai.ApiClient(self.api_config) as api_client:
@@ -152,8 +155,9 @@ def init_config(self, bv: BinaryView):
                 else:
                     self.binary_id = api_response.data.results[0].binary_id
                     self.analysis_id = api_response.data.results[0].analysis_id
+                    self.model_id = api_response.data.results[0].model_id
                     log_info(f"RevEng.AI | Binary found in RevEng.AI, binary_id: {self.binary_id}")
-                    self.set_current_info(self.binary_id, self.analysis_id)
+                    self.set_current_info(self.binary_id, self.analysis_id, self.model_id)
 
             return True, ""
 

reai_toolkit/features/match_current_function/match_current_function.py

@@ -94,7 +94,7 @@ def parse_confidence(item):
             })
 
             schema_ann_model = revengai.FunctionMatchingRequest.from_dict({
-                "model_id": 21,
+                "model_id": self.config.model_id,
                 "function_ids": [analyzed_function["function_id"]],
                 "filters": filters,
                 "result_per_function": 20,

reai_toolkit/utils/monitoring/process_binary_monitor.py

@@ -53,8 +53,16 @@ def _worker(bv: BinaryView, bid: int, aid: int):
                             f"RevEng.AI | Scheduled next status check for: {basename(bv.file.filename)} [Binary ID: {bid}] [Analysis ID: {aid}]"
                         )
                 else:
-                    callback(bid, aid)
-                    log_info(f"RevEng.AI | Analysis completed with status: {status} for Binary ID: {bid} and Analysis ID: {aid}")
+
+                    # Anaysis is complete, fetch model_id and invoke callback
+                    with revengai.ApiClient(api_config) as api_client:
+                        api_instance = revengai.AnalysesCoreApi(api_client)
+                        analysis_details: revengai.BaseResponseBasic = api_instance.get_analysis_basic_info(
+                            analysis_id=analysis_id
+                        )
+                        model_id = analysis_details.data.model_id
+                        callback(bid, aid, model_id)
+                        log_info(f"RevEng.AI | Analysis completed with status: {status} for Binary ID: {bid} | Analysis ID: {aid} | Model ID: {model_id}")
             except RequestException as ex:
                 log_error(f"RevEng.AI | Error getting binary analysis status: {str(ex)}")
             except Exception as ex: