Merge pull request #7 from RedTanny/multi_req_llm

add concurrent to executor

Author: zvigrinberg

src/concurrent_executor.py

@@ -1,9 +1,10 @@
 #!/usr/bin/env python3
 import logging
 from concurrent.futures.thread import ThreadPoolExecutor
+from typing import TYPE_CHECKING
 
-from execution import Execution
-
+if TYPE_CHECKING:
+    from execution import Execution
 
 logger = logging.getLogger(__name__)
 
@@ -39,7 +40,7 @@ class ConcurrentExecutor:
             num_of_processes: Number of concurrent requests that will be sent at once to agent, by default, it's 2.
 
     """
-    def __init__(self, execution: Execution, num_of_processes: int = 2):
+    def __init__(self, execution: 'Execution', num_of_processes: int = 2):
 
         self.execution = execution
         self.num_of_process =  num_of_processes
@@ -50,6 +51,9 @@ def execute_one_request(self, payload_data: ExecutionInput):
         :return: output object containing result from agent, or crafted output object
          in case error was intercepted on invocation.
         """
+        # Import here to avoid circular import
+        from execution import Execution
+        
         try:
             logger.info(f"Sending Request of test #{payload_data.number},"
                         f"Test-id={Execution.generate_test_id(payload_data.payload)}")

src/execution.py

@@ -14,6 +14,7 @@
 import copy
 import requests
 import time
+from concurrent_executor import ConcurrentExecutor, ExecutionInput
 logger = logging.getLogger(__name__)
 
 
@@ -423,7 +424,7 @@ def run(self) -> None:
         """
         Main execution method that orchestrates the workflow.
         
-        Reads scan configuration, generates payloads, sends requests,
+        Reads scan configuration, generates payloads, sends requests concurrently,
         and saves results for each test and iteration.
         """
         logger.info("Starting vulnerability automation execution")
@@ -433,38 +434,39 @@ def run(self) -> None:
 
             logger.info(f"Found {len(tests)} test(s) with {iterations} iteration(s) each")
 
-            # Process each test
+            # Phase 1: Collect all payloads for all test/iteration combinations
+            execution_payloads = []
             for test_idx, test in enumerate(tests, 1):
-                logger.info(f"Processing test {test_idx}/{len(tests)}: "
+                logger.info(f"Preparing test {test_idx}/{len(tests)}: "
                            f"language={test.get('language')}, "
                            f"vuln_id={test.get('vuln_id')}")
 
-                # Run iterations for this test
+                # Generate payloads for all iterations of this test
                 for iteration in range(1, iterations + 1):
-                    logger.info(f"Running iteration {iteration}/{iterations}")
-                    
                     try:
                         # Generate payload
                         payload, scan_id = self.generate_payload(test, iteration)
 
-                        # Send request
-                        result = self.send_request(payload)
-                        
-                        # Save successful result
-                        self.save_result(result, scan_id, iteration, failed=False)
-                        
-                        logger.info(f"Successfully completed iteration {iteration} "
-                                   f"for test {test_idx}")
+                        # Create unique number encoding test_idx and iteration
+                        unique_number = test_idx * 1000 + iteration
 
+                        execution_payloads.append(
+                            ExecutionInput(
+                                number=unique_number,
+                                test=test,
+                                payload=payload,
+                                request_id=scan_id
+                            )
+                        )
+                        logger.debug(f"Prepared payload for test {test_idx}, iteration {iteration}")
 
                     except Exception as e:
-                        logger.error(f"Failed iteration {iteration} for test {test_idx}: {e}")
+                        logger.error(f"Failed generating payload for test {test_idx}, "
+                                   f"iteration {iteration}: {e}")
 
-                        # Save failed result if we have a scan_id
+                        # Save failed result for payload generation error
                         try:
-                            # Generate scan_id if not already generated
-                            if 'scan_id' not in locals():
-                                scan_id = self.generate_scan_id()
+                            scan_id = self.generate_scan_id()
 
                             # Create error result
                             error_result = {
@@ -479,6 +481,43 @@ def run(self) -> None:
                         except Exception as save_error:
                             logger.error(f"Failed to save error result: {save_error}")
 
+            # Phase 2: Execute all requests concurrently
+            if execution_payloads:
+                logger.info(f"Executing {len(execution_payloads)} request(s) concurrently")
+                concurrent_executor = ConcurrentExecutor(execution=self, num_of_processes=2)
+                responses = concurrent_executor.run_requests(execution_payloads)
+                
+                # Phase 3: Process results and save them
+                for result in responses:
+                    # Extract test_idx and iteration from unique number
+                    test_idx = result.number // 1000
+                    iteration = result.number % 1000
+                    
+                    if result.intercepted_error and result.intercepted_error.strip() != "":
+                        # Request execution failed
+                        logger.error(f"Failed iteration {iteration} for test {test_idx}: "
+                                   f"{result.intercepted_error}")
+                        
+                        # Create error result
+                        error_result = {
+                            'error': result.intercepted_error,
+                            'scan_id': result.request_id,
+                            'iteration': iteration,
+                            'test': result.test
+                        }
+                        self.save_result(
+                            error_result, result.request_id, iteration, failed=True
+                        )
+                    else:
+                        # Request succeeded
+                        self.save_result(
+                            result.response, result.request_id, iteration, failed=False
+                        )
+                        logger.info(f"Successfully completed iteration {iteration} "
+                                   f"for test {test_idx}")
+            else:
+                logger.warning("No payloads to execute")
+            
             logger.info("Vulnerability automation execution completed")
 
         except FileNotFoundError as e:

src/vulnerability_main_automation.py

@@ -10,13 +10,13 @@
 import logging
 
 try:
-    from .config import Configuration, GoogleSheetsConfig, GoogleSheetsManager
+    from .config import GoogleSheetsConfig
     from .execution import Execution
     from .data_extractor import DataExtractor
     from .analysis import Analysis, GoogleSheetsDataLoader
     from .cleanup import Cleanup
 except ImportError:
-    from config import Configuration, GoogleSheetsConfig
+    from config import GoogleSheetsConfig
     from execution import Execution
     from data_extractor import DataExtractor
     from analysis import Analysis, GoogleSheetsDataLoader