diff --git a/reai_toolkit/features/match_current_function/match_current_function.py b/reai_toolkit/features/match_current_function/match_current_function.py index ed2df5d..5b73277 100755 --- a/reai_toolkit/features/match_current_function/match_current_function.py +++ b/reai_toolkit/features/match_current_function/match_current_function.py @@ -32,11 +32,11 @@ def parse_confidence(item): 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}") @@ -44,7 +44,7 @@ def parse_confidence(item): selected_collections = [int(c) for c in selected_collections_string.split(",")] if selected_binaries_string: selected_binaries = [int(b) for b in selected_binaries_string.split(",")] - + log_info(f"RevEng.AI | Selected collections: {selected_collections}") log_info(f"RevEng.AI | Selected binaries: {selected_binaries}") @@ -55,18 +55,16 @@ def parse_confidence(item): analysis_id = self.config.get_analysis_id(bv) if not analysis_id: raise Exception("Analysis not found. Please choose one using the 'Attach to existing' feature.") - + if self.cancelled.is_set(): return False, "Operation cancelled" - + with self.config.create_api_client() as api_client: analysis_core_instance = revengai.AnalysesResultsMetadataApi(api_client) analyzed_functions = analysis_core_instance.get_functions_list(analysis_id) analyzed_functions = analyzed_functions.to_dict()["data"]["functions"] if self.cancelled.is_set(): return False, "Operation cancelled" - - function_ids = [] log_info(f"RevEng.AI | Found {len(analyzed_functions)} analyzed functions") @@ -87,102 +85,121 @@ def parse_confidence(item): for func in analyzed_functions } - total_matched_functions = 0 - filters = revengai.FunctionMatchingFilters.from_dict({ - "collection_ids": selected_collections, - "binary_ids": selected_binaries + function_ids = [analyzed_function["function_id"]] + + filters = revengai.MatchFilters.from_dict({ + "collection_ids": selected_collections or None, + "binary_ids": selected_binaries or None, }) - schema_ann_model = revengai.FunctionMatchingRequest.from_dict({ - "model_id": self.config.model_id, - "function_ids": [analyzed_function["function_id"]], + start_body = revengai.StartMatchingForFunctionsInputBody.from_dict({ + "function_ids": function_ids, "filters": filters, "results_per_function": 20, - "min_similarity": similarity_threshold + "min_similarity": similarity_threshold, }) - matched_count = 0 + + with self.config.create_api_client() as api_client: + analysis_core_instance = revengai.FunctionsCoreApi(api_client) + analysis_core_instance.start_functions_matching(start_body) + + elapsed = 0.0 while True: - time.sleep(3) + if self.cancelled.is_set(): + return False, "Operation cancelled" + time.sleep(self._POLL_INTERVAL) + elapsed += self._POLL_INTERVAL + with self.config.create_api_client() as api_client: analysis_core_instance = revengai.FunctionsCoreApi(api_client) - functions_by_distance = analysis_core_instance.batch_function_matching( schema_ann_model) - - ################# - for function_by_distance in functions_by_distance.matches: - try: - if len(function_by_distance.matched_functions) == 0: - continue - log_info(f"RevEng.AI | Matched function: {function_by_distance.matched_functions[0]}") - line = { - "icon_path": f"{os.path.dirname(__file__)}/../../images/failed.png", - "icon_text": "Failed", - "source_function_id": function_by_distance.function_id, - "matched_function_name": function_by_distance.matched_functions[0].function_name, - "matched_mangled_name": function_by_distance.matched_functions[0].mangled_name, - "signature": "N/A", - "matched_hash": function_by_distance.matched_functions[0].sha_256_hash, - "matched_binary_name": function_by_distance.matched_functions[0].binary_name, - "similarity": f"{(function_by_distance.matched_functions[0].similarity):.2f}%", - "confidence": f"{(function_by_distance.matched_functions[0].confidence):.2f}%", - "nearest_neighbor_id": function_by_distance.matched_functions[0].function_id, - } - - func_addr = id_to_addr.get(function_by_distance.function_id) - log_info(f"RevEng.AI | Function ID: {function_by_distance.function_id}") - if not func_addr: - line["error"] = "Function not found in binary" - if line not in result["data"]: - result["data"].append(line) - continue - - source_function = bv.get_function_at(func_addr) - if not source_function: - source_function = bv.get_function_at(func_addr + bv.image_base) - if not source_function: - line["error"] = "Function not found in binary" - if line not in result["data"]: - result["data"].append(line) - continue - line["function_address"] = source_function.start - log_info(f"RevEng.AI | Function: {source_function.name} at {source_function.start:x}") - - if not line["matched_function_name"] or line["matched_function_name"].startswith(("sub_", "FUN_")): - line["error"] = "Function name is also debug symbol" - log_info(f"RevEng.AI | Function name is also debug symbol: {line}") - if line not in result["data"]: - result["data"].append(line) - continue - - if function_by_distance.matched_functions[0].similarity < similarity_threshold: - line["error"] = "Function score is below confidence threshold" - log_info(f"RevEng.AI | Function score is below confidence threshold: {line}") - else: - line["icon_path"] = f"{os.path.dirname(__file__)}/../../images/success.png" - line["icon_text"] = "Success" - matched_count += 1 - + status = analysis_core_instance.get_functions_matching_status(function_ids=function_ids) + + log_info(f"RevEng.AI | Matching status: {status.status} ({status.step_index}/{status.steps_total})") + + if status.status == revengai.TaskStatus.COMPLETED: + break + if status.status == revengai.TaskStatus.FAILED: + raise Exception(self._matching_error_message(status.messages)) + if elapsed >= self._POLL_TIMEOUT: + raise Exception(f"Function matching timed out after {self._POLL_TIMEOUT:.0f}s") + + if self.cancelled.is_set(): + return False, "Operation cancelled" + + with self.config.create_api_client() as api_client: + analysis_core_instance = revengai.FunctionsCoreApi(api_client) + matches_result = analysis_core_instance.get_functions_matches(function_ids=function_ids) + + matched_count = 0 + for function_by_distance in (matches_result.matches or []): + try: + if len(function_by_distance.matched_functions) == 0: + continue + log_info(f"RevEng.AI | Matched function: {function_by_distance.matched_functions[0]}") + line = { + "icon_path": f"{os.path.dirname(__file__)}/../../images/failed.png", + "icon_text": "Failed", + "source_function_id": function_by_distance.function_id, + "matched_function_name": function_by_distance.matched_functions[0].function_name, + "matched_mangled_name": function_by_distance.matched_functions[0].mangled_name, + "signature": "N/A", + "matched_hash": function_by_distance.matched_functions[0].sha_256_hash, + "matched_binary_name": function_by_distance.matched_functions[0].binary_name, + "similarity": f"{(function_by_distance.matched_functions[0].similarity):.2f}%", + "confidence": f"{(function_by_distance.matched_functions[0].confidence):.2f}%", + "nearest_neighbor_id": function_by_distance.matched_functions[0].function_id, + } + + func_addr = id_to_addr.get(function_by_distance.function_id) + log_info(f"RevEng.AI | Function ID: {function_by_distance.function_id}") + if not func_addr: + line["error"] = "Function not found in binary" + if line not in result["data"]: + result["data"].append(line) + continue + + source_function = bv.get_function_at(func_addr) + if not source_function: + source_function = bv.get_function_at(func_addr + bv.image_base) + if not source_function: + line["error"] = "Function not found in binary" + if line not in result["data"]: + result["data"].append(line) + continue + line["function_address"] = source_function.start + log_info(f"RevEng.AI | Function: {source_function.name} at {source_function.start:x}") + + if not line["matched_function_name"] or line["matched_function_name"].startswith(("sub_", "FUN_")): + line["error"] = "Function name is also debug symbol" + log_info(f"RevEng.AI | Function name is also debug symbol: {line}") if line not in result["data"]: result["data"].append(line) + continue + + if function_by_distance.matched_functions[0].similarity < similarity_threshold: + line["error"] = "Function score is below confidence threshold" + log_info(f"RevEng.AI | Function score is below confidence threshold: {line}") + else: + line["icon_path"] = f"{os.path.dirname(__file__)}/../../images/success.png" + line["icon_text"] = "Success" + matched_count += 1 + + if line not in result["data"]: + result["data"].append(line) + + except Exception as e: + log_error(f"RevEng.AI | Error processing function {function_by_distance.function_id}: {str(e)}") - except Exception as e: - log_error(f"RevEng.AI | Error processing function {function_by_distance.function_id}: {str(e)}") - #populate_table_function(result["data"]) - ################# - if functions_by_distance.status.lower() == "completed": - break - elif functions_by_distance.status.lower() == "error": - raise Exception("Function matching failed") sorted_list = sorted(result["data"], key=parse_confidence, reverse=True) result["data"] = sorted_list - #populate_table_function(result["data"]) result["failed"] = len(analyzed_functions) - matched_count result["matched"] = matched_count if self.cancelled.is_set(): return False, "Operation cancelled" - + return True, result - + except Exception as e: log_error(f"RevEng.AI | Error matching functions: {str(e)}") raise e @@ -197,14 +214,14 @@ def _process_rename_batch(self, chunk: List[Dict], bv: BinaryView, deci: Decompi try: if self.cancelled.is_set(): return 0, 0 - + addr = int(result['function_address']) if rename_function_util(self.config, bv, addr, result["matched_function_name"], result["matched_mangled_name"], result["source_function_id"]): renamed_count += 1 - + if result.get('signature_data', None) is not None: log_info(f"RevEng.AI | Applying data types for 0x{addr:x}") - + if deci is not None: log_info(f"RevEng.AI | Applying data types for 0x{addr:x} with decompiler {deci}") try: @@ -213,14 +230,14 @@ def _process_rename_batch(self, chunk: List[Dict], bv: BinaryView, deci: Decompi log_info(f"RevEng.AI | Successfully applied data types for 0x{addr:x}") except Exception as e: log_error(f"RevEng.AI | Failed to apply data types for 0x{addr:x}: {str(e)}") - - + + except (ValueError, TypeError): log_error(f"RevEng.AI | Invalid function address: {result}") continue return renamed_count, datatype_count - + except Exception as e: log_error(f"RevEng.AI | Error processing rename batch: {str(e)}") return 0, 0 @@ -239,25 +256,24 @@ def rename_functions(self, bv: BinaryView, selected_results: List[Dict]) -> List with ThreadPoolExecutor(max_workers=4) as executor: future_to_chunk = { - executor.submit(self._process_rename_batch, chunk, bv, deci): i + executor.submit(self._process_rename_batch, chunk, bv, deci): i for i, chunk in enumerate(chunks) } for future in as_completed(future_to_chunk): chunk_index = future_to_chunk[future] - try: + try: renamed_count, datatype_count = future.result() total_renamed_count += renamed_count log_info(f"RevEng.AI | Chunk {chunk_index} completed: renamed {renamed_count} functions, applied {datatype_count} data types") except Exception as e: log_error(f"RevEng.AI | Error processing chunk {chunk_index}: {str(e)}") - + success_message = f"Successfully renamed {total_renamed_count} functions!" if total_renamed_count > 0 else "No functions were renamed!" - + log_info(f"RevEng.AI | {success_message}") return True, success_message except Exception as e: log_error(f"RevEng.AI | Error renaming functions: {str(e)}") return False, str(e) - \ No newline at end of file diff --git a/reai_toolkit/features/match_functions/match_functions.py b/reai_toolkit/features/match_functions/match_functions.py index fceb485..2f8ff58 100755 --- a/reai_toolkit/features/match_functions/match_functions.py +++ b/reai_toolkit/features/match_functions/match_functions.py @@ -35,7 +35,7 @@ def parse_confidence(item): selected_collections = [int(c) for c in selected_collections_string.split(",")] if selected_binaries_string: selected_binaries = [int(b) for b in selected_binaries_string.split(",")] - + log_info(f"RevEng.AI | Selected collections: {selected_collections}") log_info(f"RevEng.AI | Selected binaries: {selected_binaries}") @@ -45,17 +45,17 @@ def parse_confidence(item): analysis_id = self.config.get_analysis_id(bv) if not analysis_id: raise Exception("Analysis not found. Please choose one using the 'Attach to existing' feature.") - + if self.cancelled.is_set(): return False, "Operation cancelled" - + with self.config.create_api_client() as api_client: analysis_core_instance = revengai.AnalysesResultsMetadataApi(api_client) analyzed_functions = analysis_core_instance.get_functions_list(analysis_id) analyzed_functions = analyzed_functions.to_dict()["data"]["functions"] if self.cancelled.is_set(): return False, "Operation cancelled" - + function_ids = [] log_info(f"RevEng.AI | Found {len(analyzed_functions)} analyzed functions") @@ -66,156 +66,142 @@ def parse_confidence(item): log_info(f"RevEng.AI | Found {len_functions} functions and {len(analyzed_functions)} analyzed functions.") for index, function in enumerate(functions, 1): - #log_info( f"RevEng.AI | Searching for {function.name} [{index}/{len_functions}]") if self.cancelled.is_set(): return False, "Operation cancelled" - + analyzed_function = next((f for f in analyzed_functions if f["function_vaddr"] == function.start), None) if analyzed_function: - #log_info(f"RevEng.AI | Found function {function.name} at {function.start:x}") function_ids.append(analyzed_function["function_id"]) else: - result["skipped"] += 1 - """ - result["data"].append({ - "icon_text": "Failed", - "function_address": function.start, - "function_name": function.name, - "demangled_name": "N/A", - "matched_name": "N/A", - "signature": "N/A", - "matched_binary": "N/A", - "similarity": "0.0%", - "confidence": "0.0%", - "error": "No Similar Function Found", - "function_address": function.start - }) - """ - - chunk_size = 50 - chunks = [function_ids[i:i + chunk_size] for i in range(0, len(function_ids), chunk_size)] - - log_info(f"RevEng.AI | Processing {len(function_ids)} functions in {len(chunks)} chunks of size {chunk_size}") + result["skipped"] += 1 id_to_addr = { func["function_id"]: func["function_vaddr"] for func in analyzed_functions } - total_matched_functions = 0 - - filters = revengai.FunctionMatchingFilters.from_dict({ - "collection_ids": selected_collections, - "binary_ids": selected_binaries + log_info(f"RevEng.AI | Matching {len(function_ids)} functions in analysis {analysis_id}") + + filters = revengai.MatchFilters.from_dict({ + "collection_ids": selected_collections or None, + "binary_ids": selected_binaries or None, }) - - schema_ann_model = revengai.AnalysisFunctionMatchingRequest.from_dict({ - "function_ids": function_ids, + + start_body = revengai.StartMatchingForAnalysisInputBody.from_dict({ "min_similarity": 0, "filters": filters, - "results_per_function": 1 + "results_per_function": 1, }) - matched_count = 0 + + with self.config.create_api_client() as api_client: + analyses_core_instance = revengai.AnalysesCoreApi(api_client) + analyses_core_instance.start_analysis_function_matching(analysis_id, start_body) + + elapsed = 0.0 while True: - time.sleep(3) + if self.cancelled.is_set(): + return False, "Operation cancelled" + time.sleep(self._POLL_INTERVAL) + elapsed += self._POLL_INTERVAL + with self.config.create_api_client() as api_client: - analysis_core_instance = revengai.FunctionsCoreApi(api_client) - functions_by_distance = analysis_core_instance.analysis_function_matching(analysis_id, schema_ann_model) + analyses_core_instance = revengai.AnalysesCoreApi(api_client) + status = analyses_core_instance.get_analysis_function_matching_status(analysis_id) - ################# - if self.cancelled.is_set(): - return False, "Operation cancelled" - - for function_by_distance in functions_by_distance.matches: - try: - if len(function_by_distance.matched_functions) == 0: - continue - log_info(f"RevEng.AI | Matched function: {function_by_distance.matched_functions[0]}") - line = { - "icon_path": f"{os.path.dirname(__file__)}/../../images/failed.png", - "icon_text": "Failed", - "function_address": function_by_distance.matched_functions[0].function_vaddr, - "function_name": "N/A", - "source_function_id": function_by_distance.function_id, - "matched_function_name": function_by_distance.matched_functions[0].function_name, - "matched_mangled_name": function_by_distance.matched_functions[0].mangled_name, - "signature": "N/A", - "matched_hash": function_by_distance.matched_functions[0].sha_256_hash, - "matched_binary_name": function_by_distance.matched_functions[0].binary_name, - "similarity": f"{(function_by_distance.matched_functions[0].similarity):.2f}%", - "confidence": f"{(function_by_distance.matched_functions[0].confidence):.2f}%", - "error": "", - "nearest_neighbor_id": function_by_distance.matched_functions[0].function_id, - "function_address": "N/A" - } - - func_addr = id_to_addr.get(function_by_distance.function_id) - log_info(f"RevEng.AI | Function address: {func_addr:x}") - log_info(f"RevEng.AI | Function ID: {function_by_distance.function_id}") - if not func_addr: - line["error"] = "Function not found in binary" - if line not in result["data"]: - result["data"].append(line) - continue - - function = bv.get_function_at(func_addr) - log_info(f"RevEng.AI | Function: {function} at {func_addr:x}") - if function: - line["function_name"] = function.name - line["function_address"] = function.start - if not function: - function = bv.get_function_at(func_addr + bv.image_base) - log_info(f"RevEng.AI | Function: {function} at {func_addr + bv.image_base:x}") - if function: - line["function_name"] = function.name - line["function_address"] = function.start - else: - """ - line["error"] = "Function not found in binary" - if line not in result["data"]: - result["data"].append(line) - """ - continue - - if not line["matched_function_name"] or line["matched_function_name"].startswith(("sub_", "FUN_")): - line["error"] = "Function name is also debug symbol" - log_info(f"RevEng.AI | Function name is also debug symbol: {line}") - if line not in result["data"]: - result["data"].append(line) - continue - - if function_by_distance.matched_functions[0].confidence < confidence_threshold: - line["error"] = "Function score is below confidence threshold" - log_info(f"RevEng.AI | Function score is below confidence threshold: {line}") - else: - line["icon_path"] = f"{os.path.dirname(__file__)}/../../images/success.png" - line["icon_text"] = "Success" - matched_count += 1 - - if line not in result["data"] and line["function_address"] != "N/A": + log_info(f"RevEng.AI | Matching status: {status.status} ({status.step_index}/{status.steps_total})") + + if status.status == revengai.TaskStatus.COMPLETED: + break + if status.status == revengai.TaskStatus.FAILED: + raise Exception(self._matching_error_message(status.messages)) + if elapsed >= self._POLL_TIMEOUT: + raise Exception(f"Function matching timed out after {self._POLL_TIMEOUT:.0f}s") + + if self.cancelled.is_set(): + return False, "Operation cancelled" + + with self.config.create_api_client() as api_client: + analyses_core_instance = revengai.AnalysesCoreApi(api_client) + matches_result = analyses_core_instance.get_analysis_function_matches(analysis_id) + + matched_count = 0 + for function_by_distance in (matches_result.matches or []): + try: + if len(function_by_distance.matched_functions) == 0: + continue + log_info(f"RevEng.AI | Matched function: {function_by_distance.matched_functions[0]}") + line = { + "icon_path": f"{os.path.dirname(__file__)}/../../images/failed.png", + "icon_text": "Failed", + "function_address": function_by_distance.matched_functions[0].function_vaddr, + "function_name": "N/A", + "source_function_id": function_by_distance.function_id, + "matched_function_name": function_by_distance.matched_functions[0].function_name, + "matched_mangled_name": function_by_distance.matched_functions[0].mangled_name, + "signature": "N/A", + "matched_hash": function_by_distance.matched_functions[0].sha_256_hash, + "matched_binary_name": function_by_distance.matched_functions[0].binary_name, + "similarity": f"{(function_by_distance.matched_functions[0].similarity):.2f}%", + "confidence": f"{(function_by_distance.matched_functions[0].confidence):.2f}%", + "error": "", + "nearest_neighbor_id": function_by_distance.matched_functions[0].function_id, + "function_address": "N/A" + } + + func_addr = id_to_addr.get(function_by_distance.function_id) + log_info(f"RevEng.AI | Function ID: {function_by_distance.function_id}") + if not func_addr: + line["error"] = "Function not found in binary" + if line not in result["data"]: + result["data"].append(line) + continue + + function = bv.get_function_at(func_addr) + log_info(f"RevEng.AI | Function: {function} at {func_addr:x}") + if function: + line["function_name"] = function.name + line["function_address"] = function.start + if not function: + function = bv.get_function_at(func_addr + bv.image_base) + log_info(f"RevEng.AI | Function: {function} at {func_addr + bv.image_base:x}") + if function: + line["function_name"] = function.name + line["function_address"] = function.start + else: + continue + + if not line["matched_function_name"] or line["matched_function_name"].startswith(("sub_", "FUN_")): + line["error"] = "Function name is also debug symbol" + log_info(f"RevEng.AI | Function name is also debug symbol: {line}") + if line not in result["data"]: result["data"].append(line) + continue + + if function_by_distance.matched_functions[0].confidence < confidence_threshold: + line["error"] = "Function score is below confidence threshold" + log_info(f"RevEng.AI | Function score is below confidence threshold: {line}") + else: + line["icon_path"] = f"{os.path.dirname(__file__)}/../../images/success.png" + line["icon_text"] = "Success" + matched_count += 1 + + if line not in result["data"] and line["function_address"] != "N/A": + result["data"].append(line) + + except Exception as e: + log_error(f"RevEng.AI | Error processing function {function_by_distance.function_id}: {str(e)}") - except Exception as e: - log_error(f"RevEng.AI | Error processing function {function_by_distance.function_id}: {str(e)}") - #populate_table_function(result["data"]) - ################# - if functions_by_distance.status.lower() == "completed": - - break - elif functions_by_distance.status.lower() == "error": - raise Exception("Function matching failed") sorted_list = sorted(result["data"], key=parse_confidence, reverse=True) result["data"] = sorted_list - #populate_table_function(result["data"]) result["failed"] = len(analyzed_functions) - matched_count result["matched"] = matched_count if self.cancelled.is_set(): return False, "Operation cancelled" - + return True, result - + except Exception as e: log_error(f"RevEng.AI | Error matching functions: {str(e)}") raise e @@ -230,14 +216,14 @@ def _process_rename_batch(self, config, chunk: List[Dict], bv: BinaryView, deci: try: if self.cancelled.is_set(): return 0, 0 - + addr = int(result['function_address']) if rename_function_util(config, bv, addr, result["matched_function_name"], result["matched_mangled_name"], result["source_function_id"]): renamed_count += 1 - + if result.get('signature_data', None) is not None: log_info(f"RevEng.AI | Applying data types for 0x{addr:x}") - + if deci is not None: log_info(f"RevEng.AI | Applying data types for 0x{addr:x} with decompiler {deci}") try: @@ -246,14 +232,14 @@ def _process_rename_batch(self, config, chunk: List[Dict], bv: BinaryView, deci: log_info(f"RevEng.AI | Successfully applied data types for 0x{addr:x}") except Exception as e: log_error(f"RevEng.AI | Failed to apply data types for 0x{addr:x}: {str(e)}") - - + + except (ValueError, TypeError): log_error(f"RevEng.AI | Invalid function address: {result}") continue return renamed_count, datatype_count - + except Exception as e: log_error(f"RevEng.AI | Error processing rename batch: {str(e)}") return 0, 0 @@ -272,25 +258,24 @@ def rename_functions(self, bv: BinaryView, selected_results: List[Dict]) -> List with ThreadPoolExecutor(max_workers=4) as executor: future_to_chunk = { - executor.submit(self._process_rename_batch, self.config, chunk, bv, deci): i + executor.submit(self._process_rename_batch, self.config, chunk, bv, deci): i for i, chunk in enumerate(chunks) } for future in as_completed(future_to_chunk): chunk_index = future_to_chunk[future] - try: + try: renamed_count, datatype_count = future.result() total_renamed_count += renamed_count log_info(f"RevEng.AI | Chunk {chunk_index} completed: renamed {renamed_count} functions, applied {datatype_count} data types") except Exception as e: log_error(f"RevEng.AI | Error processing chunk {chunk_index}: {str(e)}") - + success_message = f"Successfully renamed {total_renamed_count} functions!" if total_renamed_count > 0 else "No functions were renamed!" - + log_info(f"RevEng.AI | {success_message}") return True, success_message except Exception as e: log_error(f"RevEng.AI | Error renaming functions: {str(e)}") return False, str(e) - \ No newline at end of file diff --git a/reai_toolkit/utils/features/matching.py b/reai_toolkit/utils/features/matching.py index a20c120..3a212d7 100755 --- a/reai_toolkit/utils/features/matching.py +++ b/reai_toolkit/utils/features/matching.py @@ -9,6 +9,9 @@ from concurrent.futures import ThreadPoolExecutor, as_completed class MatchFeature: + _POLL_INTERVAL = 3.0 + _POLL_TIMEOUT = 1200.0 + def __init__(self, config): self.config = config self.base_addr = None @@ -27,7 +30,15 @@ def cancel(self): def clear_cancelled(self): log_info("RevEng.AI | Clearing cancelled event...") self.cancelled.clear() - + + @staticmethod + def _matching_error_message(messages) -> str: + """Build a human-readable error from a matching workflow's progress messages.""" + if not messages: + return "Function matching failed." + errors = [m.text for m in messages if (getattr(m, "level", "") or "").upper() == "ERROR"] + return "; ".join(errors or [messages[-1].text]) + # Search collections/Binaries Process Functions def search_items(self, bv: BinaryView, options: Dict[str, Any]): diff --git a/tests/headless/test_match_current_function.py b/tests/headless/test_match_current_function.py index 38c0bab..e710545 100644 --- a/tests/headless/test_match_current_function.py +++ b/tests/headless/test_match_current_function.py @@ -45,9 +45,10 @@ def test_match_current_function_pipeline(bv, mocker): by_distance.function_id = 100 by_distance.matched_functions = [matched] core_api = mocker.patch.object(mcf_mod.revengai, "FunctionsCoreApi").return_value - core_api.batch_function_matching.return_value = MagicMock( - status="completed", matches=[by_distance] + core_api.get_functions_matching_status.return_value = MagicMock( + status="COMPLETED", step_index=1, steps_total=1 ) + core_api.get_functions_matches.return_value = MagicMock(matches=[by_distance]) ok, result = feature.match_functions( bv, {"function": func.start, "similarity_threshold": 90} diff --git a/tests/unit/sdk/test_sdk_schemas.py b/tests/unit/sdk/test_sdk_schemas.py index 3c27544..b4f2de1 100644 --- a/tests/unit/sdk/test_sdk_schemas.py +++ b/tests/unit/sdk/test_sdk_schemas.py @@ -34,10 +34,18 @@ "get_analysis_status", "get_analysis_basic_info", "get_analysis_function_map", + "start_analysis_function_matching", + "get_analysis_function_matching_status", + "get_analysis_function_matches", ], "ModelsApi": ["get_models"], "AnalysesResultsMetadataApi": ["get_functions_list"], - "FunctionsCoreApi": ["analysis_function_matching", "auto_unstrip"], + "FunctionsCoreApi": [ + "start_functions_matching", + "get_functions_matching_status", + "get_functions_matches", + "auto_unstrip", + ], "FunctionsRenamingHistoryApi": ["rename_function_id"], "FunctionsDataTypesApi": [ "generate_function_data_types_for_functions", @@ -120,9 +128,24 @@ def test_analysis_id_methods_accept_analysis_id(): assert "analysis_id" in _params(getattr(revengai.AnalysesCoreApi, method)) -def test_function_matching_accepts_request_kwargs(): - params = _params(revengai.FunctionsCoreApi.analysis_function_matching) - assert {"analysis_id", "analysis_function_matching_request"} <= params +def test_start_functions_matching_accepts_request_kwargs(): + params = _params(revengai.FunctionsCoreApi.start_functions_matching) + assert "start_matching_for_functions_input_body" in params + + +def test_functions_matching_status_and_results_accept_function_ids(): + for method in ("get_functions_matching_status", "get_functions_matches"): + assert "function_ids" in _params(getattr(revengai.FunctionsCoreApi, method)) + + +def test_start_analysis_function_matching_accepts_request_kwargs(): + params = _params(revengai.AnalysesCoreApi.start_analysis_function_matching) + assert {"analysis_id", "start_matching_for_analysis_input_body"} <= params + + +def test_analysis_matching_status_and_results_accept_analysis_id(): + for method in ("get_analysis_function_matching_status", "get_analysis_function_matches"): + assert "analysis_id" in _params(getattr(revengai.AnalysesCoreApi, method)) def test_auto_unstrip_accepts_request_kwargs(): @@ -157,9 +180,9 @@ def test_analysis_create_request_has_plugin_fields(): } <= set(revengai.AnalysisCreateRequest.model_fields) -def test_function_matching_filters_has_plugin_fields(): - assert {"collection_ids", "binary_ids"} <= set( - revengai.FunctionMatchingFilters.model_fields +def test_match_filters_has_plugin_fields(): + assert {"collection_ids", "binary_ids", "debug_types"} <= set( + revengai.MatchFilters.model_fields ) @@ -275,33 +298,35 @@ def test_inline_comment_has_plugin_fields(): assert {"line", "comment"} <= set(InlineComment.model_fields) -def test_matching_request_per_function_count_field_name(): - assert "results_per_function" in revengai.AnalysisFunctionMatchingRequest.model_fields - assert "result_per_function" not in revengai.AnalysisFunctionMatchingRequest.model_fields +def test_start_matching_input_bodies_have_per_function_count_field(): + for model in ( + revengai.StartMatchingForAnalysisInputBody, + revengai.StartMatchingForFunctionsInputBody, + ): + assert "results_per_function" in model.model_fields + assert "result_per_function" not in model.model_fields -def test_matching_request_from_dict_honors_results_per_function(): - request = revengai.AnalysisFunctionMatchingRequest.from_dict( - { - "min_similarity": 0, - "filters": {"collection_ids": [], "binary_ids": []}, - "results_per_function": 7, - } - ) - assert request.to_dict().get("results_per_function") == 7 +def test_start_matching_for_functions_input_has_function_ids_field(): + assert "function_ids" in revengai.StartMatchingForFunctionsInputBody.model_fields + +def test_get_matches_output_has_status_and_matches(): + assert {"status", "matches"} <= set(revengai.GetMatchesOutputBody.model_fields) -def test_function_matching_request_has_function_ids_field(): - assert "function_ids" in revengai.FunctionMatchingRequest.model_fields + +def test_get_matches_status_output_has_progress_fields(): + assert {"status", "step", "step_index", "steps_total", "messages"} <= set( + revengai.GetMatchesStatusOutputBody.model_fields + ) def test_match_functions_request_shape_is_honored(): - filters = revengai.FunctionMatchingFilters.from_dict( + filters = revengai.MatchFilters.from_dict( {"collection_ids": [1], "binary_ids": [2]} ) - request = revengai.AnalysisFunctionMatchingRequest.from_dict( + request = revengai.StartMatchingForAnalysisInputBody.from_dict( { - "function_ids": [1, 2, 3], "min_similarity": 0, "filters": filters, "results_per_function": 1, @@ -314,12 +339,11 @@ def test_match_functions_request_shape_is_honored(): def test_match_current_function_request_shape_is_honored(): - filters = revengai.FunctionMatchingFilters.from_dict( + filters = revengai.MatchFilters.from_dict( {"collection_ids": [1], "binary_ids": [2]} ) - request = revengai.FunctionMatchingRequest.from_dict( + request = revengai.StartMatchingForFunctionsInputBody.from_dict( { - "model_id": 7, "function_ids": [42], "filters": filters, "results_per_function": 20,