dev: signatures done

Author: anisiosts

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

revengai/features/match_current_function/match_current_function.py

@@ -342,104 +342,100 @@ def rename_function(self, bv: BinaryView, selected_result: Dict) -> List[Dict]:
 
     def _process_data_type_batch(self, chunk: List[Dict]) -> List[Dict]:
         try:
-            log_info(f"RevEng.AI | Processing data types batch of {len(chunk)} items")
-            
-            nearest_neighbor_ids = [item["nearest_neighbor_id"] for item in chunk]
-            
-            response = RE_functions_data_types(nearest_neighbor_ids)
-            
-            if response.status_code != 200:
-                log_error(f"RevEng.AI | Data types API call failed with status {response.status_code}")
-                return []
-            
-            data = response.json()
-            
-            if "status" in data and data["status"] == "processing":
-                poll_id = data.get("poll_id")
-                if poll_id:
-                    log_info(f"RevEng.AI | Polling for data types with ID: {poll_id}")
-                    
-                    max_attempts = 30
-                    for attempt in range(max_attempts):
-                        time.sleep(2)
-                        poll_response = RE_functions_data_types_poll(poll_id)
-                        
-                        if poll_response.status_code == 200:
-                            poll_data = poll_response.json()
-                            if poll_data.get("status") == "completed":
-                                data = poll_data
-                                break
-                        else:
-                            log_error(f"RevEng.AI | Polling failed with status {poll_response.status_code}")
-                            break
-                    else:
-                        log_error(f"RevEng.AI | Polling timed out after {max_attempts} attempts")
-                        return []
-            
+            log_info(f"RevEng.AI | Processing chunk of {len(chunk)} functions")
+            function_ids = set([result['nearest_neighbor_id'] for result in chunk])
+            RE_functions_data_types(function_ids=list(function_ids))
             signatures = []
-            for item in data.get("data", []):
-                signature = self.make_signature(item.get("data_types", []))
-                signatures.append({
-                    "nearest_neighbor_id": item["nearest_neighbor_id"],
-                    "signature": signature
-                })
-            
+            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 | {pending_count} items still pending... trying again")
+                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 chunk:
+                    if result['nearest_neighbor_id'] == item['function_id']:
+                        signature = self.make_signature(item['data_types'])
+                        if signature != "N/A":
+                            signatures.append({"nearest_neighbor_id": result['nearest_neighbor_id'], "signature": signature})
+                        break
+
+            #log_info(f"RevEng.AI | Total count: {total_count}")
+            #log_info(f"RevEng.AI | Total data types: {total_data_types}")
+            #log_info(f"RevEng.AI | Items: {items}")
+
             return signatures
-            
         except Exception as e:
-            log_error(f"RevEng.AI | Error processing data types batch: {str(e)}")
+            log_error(f"RevEng.AI | Error processing data type batch: {str(e)}")
             return []
-
+        
     def make_signature(self, data_types: List[Dict]) -> str:
         try:
-            if not data_types:
-                return "void function();"
-            
-            # For now, create a simple signature
-            # This would need to be enhanced based on actual data_types structure
-            return_type = "void"
-            params = []
-            
-            for dt in data_types:
-                if dt.get("type") == "return":
-                    return_type = dt.get("name", "void")
-                elif dt.get("type") == "parameter":
-                    param_type = dt.get("name", "int")
-                    param_name = dt.get("param_name", f"param{len(params)}")
-                    params.append(f"{param_type} {param_name}")
-            
-            params_str = ", ".join(params) if params else ""
-            return f"{return_type} function({params_str});"
-            
+            #log_info(f"RevEng.AI | Making signature for {data_types}")
+            signature = "("
+            for _, arg in data_types['func_types'].get('header', {}).get('args', {}).items():
+                #log_info(f"RevEng.AI | Arg: {arg}")
+                signature += f"{arg.get('type', 'N/A')}, "
+            signature = signature[:-2] if signature.endswith(", ") else signature
+
+            signature += f") {data_types['func_types'].get('type', 'N/A')}"
+
+            log_info(f"RevEng.AI | Signature: {signature}")
+            return signature
         except Exception as e:
-            log_error(f"RevEng.AI | Error creating signature: {str(e)}")
-            return "void function();"
+            log_error(f"RevEng.AI | Error making signature: {str(e)}")
+            return "N/A"
 
     def fetch_data_types(self, bv: BinaryView, selected_results: List[Dict]) -> Tuple[bool, Dict[str, Any]]:
-        """Fetch data types for selected function matches"""
         try:
-            log_info(f"RevEng.AI | Starting data type fetching for {len(selected_results)} functions")
+            log_info("RevEng.AI | Starting data type fetching")
 
-            # Process in chunks to avoid API limits
+            if len(selected_results) == 0:
+                return False, "No valid functions selected"
+
             chunk_size = 50
-            all_signatures = []
-            
-            for i in range(0, len(selected_results), chunk_size):
-                chunk = selected_results[i:i+chunk_size]
-                log_info(f"RevEng.AI | Processing chunk {i//chunk_size + 1}/{(len(selected_results) + chunk_size - 1)//chunk_size}")
-                
-                signatures = self._process_data_type_batch(chunk)
-                all_signatures.extend(signatures)
-            
-            success_count = len([s for s in all_signatures if s.get("signature") != "void function();"])
-            
-            log_info(f"RevEng.AI | Data type fetching completed. {success_count} functions have signatures")
+            if len(selected_results) < chunk_size:
+                chunks = [selected_results]
+            else:
+                chunks = [selected_results[i:i + chunk_size] for i in range(0, len(selected_results), chunk_size)]
+
+            log_info(f"RevEng.AI | Processing {len(selected_results)} functions in {len(chunks)} chunks of size {chunk_size}")
+
+            signatures = []
 
-            return True, {
-                "signatures": all_signatures,
-                "success_count": success_count
+            with ThreadPoolExecutor(max_workers=4) as executor:
+                future_to_chunk = {
+                    executor.submit(self._process_data_type_batch, chunk): i
+                    for i, chunk in enumerate(chunks)
+                }
+
+                for future in as_completed(future_to_chunk):
+                    chunk_index = future_to_chunk[future]
+                    try:
+                        chunk = future.result()
+                        log_info(f"RevEng.AI | Chunk {chunk_index} completed")
+                        signatures.extend(chunk)
+        
+                    except Exception as e:
+                        log_error(f"RevEng.AI | Error processing chunk {chunk_index}: {str(e)}")
+
+            options = {
+                "success_count": len(signatures),
+                "signatures": signatures
             }
-            
+
+            return True, options
         except Exception as e:
-            log_error(f"RevEng.AI | Error in data type fetching: {str(e)}")
-            return False, str(e) 
+            log_error(f"RevEng.AI | Error fetching data types: {str(e)}")
+            return False, str(e)

revengai/features/match_current_function/match_current_function_dialog.py

@@ -305,7 +305,7 @@ def on_fetching_data_types_finished(self, success, data):
         if success:
             log_info(f"RevEng.AI | Data type fetching completed with {data['success_count']} functions having signatures")
             self.results_tab.update_current_matches_with_signatures(data["signatures"])
-            self.results_tab.populate_results_table()
+            #self.results_tab.populate_results_table()
             self.status_label.setText(f"Data type fetching completed: {data['success_count']} functions have signatures")
 
             QMessageBox.information(

revengai/features/match_functions/match_functions.py

@@ -371,31 +371,33 @@ def _process_data_type_batch(self, chunk: List[Dict]) -> List[Dict]:
             function_ids = set([result['nearest_neighbor_id'] for result in chunk])
             RE_functions_data_types(function_ids=list(function_ids))
             signatures = []
+            items = []
             while True:
                 response = RE_functions_data_types_poll(    
                     function_ids=list(function_ids),
                 ).json()
                 data = response.get("data", {})
-                total_count = data.get("total_count", 0)
-                total_data_types = data.get("total_data_types_count", 0)
                 items = data.get("items", [])
-                log_info(f"RevEng.AI | Response: {response}")
-                if not any(item.get("status") == "pending" for item in items):
+                pending_count = sum(1 for item in items if item.get("status") == "pending")
+                log_info(f"RevEng.AI | {pending_count} items still pending... trying again")
+                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 chunk:
                     if result['nearest_neighbor_id'] == item['function_id']:
                         signature = self.make_signature(item['data_types'])
                         if signature != "N/A":
                             signatures.append({"nearest_neighbor_id": result['nearest_neighbor_id'], "signature": signature})
                         break
 
-            log_info(f"RevEng.AI | Total count: {total_count}")
-            log_info(f"RevEng.AI | Total data types: {total_data_types}")
-            log_info(f"RevEng.AI | Items: {items}")
+            #log_info(f"RevEng.AI | Total count: {total_count}")
+            #log_info(f"RevEng.AI | Total data types: {total_data_types}")
+            #log_info(f"RevEng.AI | Items: {items}")
 
             return signatures
         except Exception as e:
@@ -404,12 +406,14 @@ def _process_data_type_batch(self, chunk: List[Dict]) -> List[Dict]:
 
     def make_signature(self, data_types: List[Dict]) -> str:
         try:
-            log_info(f"RevEng.AI | Making signature for {data_types}")
-            signature = ""
-            signature += f"{data_types['func_types'].get('type', 'N/A')} "
-
-            for arg in data_types['func_types'].get('args', []):
+            #log_info(f"RevEng.AI | Making signature for {data_types}")
+            signature = "("
+            for _, arg in data_types['func_types'].get('header', {}).get('args', {}).items():
+                #log_info(f"RevEng.AI | Arg: {arg}")
                 signature += f"{arg.get('type', 'N/A')}, "
+            signature = signature[:-2] if signature.endswith(", ") else signature
+
+            signature += f") {data_types['func_types'].get('type', 'N/A')}"
 
             log_info(f"RevEng.AI | Signature: {signature}")
             return signature
@@ -425,7 +429,10 @@ def fetch_data_types(self, bv: BinaryView, selected_results: List[Dict]) -> Tupl
                 return False, "No valid functions selected"
 
             chunk_size = 50
-            chunks = [selected_results[i:i + chunk_size] for i in range(0, len(selected_results), chunk_size)]
+            if len(selected_results) < chunk_size:
+                chunks = [selected_results]
+            else:
+                chunks = [selected_results[i:i + chunk_size] for i in range(0, len(selected_results), chunk_size)]
 
             log_info(f"RevEng.AI | Processing {len(selected_results)} functions in {len(chunks)} chunks of size {chunk_size}")
 

revengai/features/match_functions/match_functions_dialog.py

@@ -230,7 +230,7 @@ def on_fetching_data_types_finished(self, success, data):
         if success:
             log_info(f"RevEng.AI | Data type fetching completed with {data['success_count']} functions having signatures")
             self.results_tab.update_current_matches_with_signatures(data["signatures"])
-            self.results_tab.populate_results_table()
+            #self.results_tab.populate_results_table()
             self.status_label.setText(f"Data type fetching completed!")
             self.results_tab.status_label.setText(f"Data type fetching completed: {data['success_count']} functions have signatures")
             QMessageBox.information(self, "RevEng.AI Fetch Data Types", f"Data types fetched successfully.\n{data['success_count']} function{'' if data['success_count'] == 1 else 's'} have signatures.", QMessageBox.Ok)

revengai/features/match_functions/tab_result.py

@@ -71,6 +71,7 @@ def filter_results(self, text):
     def populate_results_table(self):
         self.selected_results.clear()
 
+        self.results_table.setRowCount(0)
         self.results_table.setRowCount(len(self.current_matches))
 
         for row, match in enumerate(self.current_matches):
@@ -94,19 +95,25 @@ def populate_results_table(self):
 
             for column, field in enumerate(column_data, start=1):
                 value = match.get(field, "N/A")
-                item = QTableWidgetItem(value if len(value) < 25 else value[:22] + "...")
+                if field != "signature":
+                    value = value if len(value) < 25 else value[:22] + "..."
+                item = QTableWidgetItem(value)
                 item.setFlags(item.flags() & ~Qt.ItemIsEditable)
                 self.results_table.setItem(row, column, item)
 
             if icon_text == "Success":
                 self.selected_results.append(match)
 
     def update_current_matches_with_signatures(self, selected_results):
+        log_info(f"RevEng.AI | Updating current matches with signatures")
         for match in self.current_matches:
-            if match.get("nearest_neighbor_id", False):
+            if not match.get("nearest_neighbor_id", False):
                 continue
             for result in selected_results:
+                if not result.get("nearest_neighbor_id", False):
+                    continue
                 if match["nearest_neighbor_id"] == result["nearest_neighbor_id"]:
+                    log_info(f"RevEng.AI | Found signature for {match['original_name']}")
                     match["signature"] = result["signature"]
                     break
         self.populate_results_table()