RedTanny/config_threads

Config threads

Author: RedTanny

src/config.py

@@ -36,6 +36,7 @@ class Configuration:
     DEFAULT_LOG_LEVEL = logging.DEBUG
     DEFAULT_LANGUAGE = None  # None means run all languages
     DEFAULT_STRICT_MODE = False  # None means run all languages
+    DEFAULT_NUM_THREADS = 2
 
     # Environment variable names
     ENV_SERVICE_URL = "SERVICE_URL"
@@ -45,6 +46,7 @@ class Configuration:
     ENV_LOG_LEVEL = "LOG_LEVEL"
     ENV_LANGUAGE = "LANGUAGE"
     ENV_STRICT_MODE = "STRICT_MODE"
+    ENV_NUM_THREADS = "NUM_THREADS"
 
     def __init__(
         self,
@@ -54,7 +56,8 @@ def __init__(
         timeout: Optional[int] = None,
         log_level: Optional[int] = None,
         language: Optional[str] = None,
-        strict_mode: Optional[str] = None
+        strict_mode: Optional[str] = None,
+        num_threads: Optional[int] = None
     ):
         """
         Initialize Configuration with provided values or defaults.
@@ -67,6 +70,7 @@ def __init__(
             log_level: Logging level
             language: Language filter (e.g., 'c', 'go', 'python'). None means run all languages.
             strict_mode: Stricter matching mode for integration tests
+            num_threads: Number of threads to use
         """
         # Resolve paths to absolute paths
         self.input_dir = os.path.abspath(
@@ -122,6 +126,18 @@ def __init__(
             self.DEFAULT_STRICT_MODE
         )
 
+        # Get num_threads with precedence
+        num_threads_str = self._get_env(self.ENV_NUM_THREADS)
+        if num_threads is not None:
+            self.num_threads = num_threads
+        elif num_threads_str:
+            try:
+                self.num_threads = int(num_threads_str)
+            except ValueError:
+                self.num_threads = self.DEFAULT_NUM_THREADS
+        else:
+            self.num_threads = self.DEFAULT_NUM_THREADS
+
 
     @classmethod
     def create(
@@ -169,6 +185,9 @@ def create(
         if hasattr(args, 'strict_mode') and args.strict_mode is not None:
             kwargs['strict_mode'] = args.strict_mode
 
+        if hasattr(args, 'num_threads') and args.num_threads is not None:
+            kwargs['num_threads'] = args.num_threads
+
 
         # Create with only the provided CLI args
         # __init__ will handle precedence: provided value > env var > default
@@ -244,7 +263,8 @@ def __repr__(self) -> str:
             f"service_url={self.service_url!r}, "
             f"timeout={self.timeout}, "
             f"log_level={self.log_level}, "
-            f"language={self.language!r})"
+            f"language={self.language!r}, "
+            f"num_threads={self.num_threads})"
         )
 
 
@@ -275,6 +295,7 @@ def __init__(
         timeout: Optional[int] = None,
         log_level: Optional[int] = None,
         language: Optional[str] = None,
+        num_threads: Optional[int] = None,
         gsheets_mode: Optional[str] = None,
         input_sheet_id: Optional[str] = None,
         output_sheet_id: Optional[str] = None,
@@ -291,6 +312,7 @@ def __init__(
             timeout: Request timeout in seconds
             log_level: Logging level
             language: Language filter (e.g., 'c', 'go', 'python'). None means run all languages.
+            num_threads: Number of threads to use
             gsheets_mode: Google Sheets mode: 'none', 'input', 'output', or 'both'
             input_sheet_id: Google Sheets ID for reading input data
             output_sheet_id: Google Sheets ID for writing results
@@ -304,7 +326,8 @@ def __init__(
             service_url=service_url,
             timeout=timeout,
             log_level=log_level,
-            language=language
+            language=language,
+            num_threads=num_threads
         )
 
         # Google Sheets mode
@@ -420,6 +443,8 @@ def create(
             kwargs['log_level'] = Configuration._parse_log_level_static(args.log_level)
         if hasattr(args, 'language') and args.language is not None:
             kwargs['language'] = args.language
+        if hasattr(args, 'num_threads') and args.num_threads is not None:
+            kwargs['num_threads'] = args.num_threads
 
         # Extract Google Sheets kwargs
         if hasattr(args, 'gsheets_mode') and args.gsheets_mode is not None:

src/execution.py

@@ -427,6 +427,10 @@ def run(self) -> None:
         Reads scan configuration, generates payloads, sends requests concurrently,
         and saves results for each test and iteration.
         """
+
+        def valid_num_threads(num_threads: int) -> bool:
+            return num_threads > 0 and num_threads <= 5
+
         logger.info("Starting vulnerability automation execution")
 
         try:
@@ -484,7 +488,10 @@ def run(self) -> None:
             # 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)
+                if not valid_num_threads(self.config.num_threads):
+                    logger.error(f"Invalid number of threads: {self.config.num_threads} set to default of 2")
+                    self.config.num_threads = 2
+                concurrent_executor = ConcurrentExecutor(execution=self, num_of_processes=self.config.num_threads)
                 responses = concurrent_executor.run_requests(execution_payloads)
 
                 # Phase 3: Process results and save them
@@ -534,6 +541,10 @@ def run(self) -> None:
             raise
 
     def prepare_tests(self, scan_config_override: Union[str, None]) -> tuple:
+
+        def valid_language( language: str) -> bool:
+            valid_languages = ["c", "go", "python","java","javascript","rust"]
+            return language.lower() in valid_languages
         # Read scan configuration
         scan_config = self.read_scan_config(scan_config_override)
 
@@ -545,9 +556,9 @@ def prepare_tests(self, scan_config_override: Union[str, None]) -> tuple:
 
         iterations = scan_config['iterations']
         tests = scan_config['tests']
-
+        self.config.language ="None"
         # Filter tests by language if specified in config
-        if self.config.language:
+        if self.config.language and valid_language(self.config.language):
             original_count = len(tests)
             tests = [
                 test for test in tests

src/vulnerability_main_automation.py

@@ -101,6 +101,14 @@ def create_argument_parser() -> argparse.ArgumentParser:
              'If not specified, runs all languages. '
              '(default: None or LANGUAGE env var)'
     )
+
+    parser.add_argument(
+        '--num-threads',
+        type=int,
+        default=None,
+        help='Number of threads to use for the vulnerability automation '
+             '(default: 2 or NUM_THREADS env var)'
+    )
 
     # Google Sheets integration arguments
     parser.add_argument(
@@ -198,6 +206,9 @@ def main():
                 if "language" in config_override:
                     logger.info("Override Language filter: %s with value %s from Google Sheets", config.language, config_override["language"])
                     config.language = config_override["language"]
+                if "num_threads" in config_override:
+                    logger.info("Override Number of threads: %s with value %s from Google Sheets", config.num_threads, config_override["num_threads"])
+                    config.num_threads = int(config_override["num_threads"])
                 logger.info("Successfully generated scan.json from Google Sheets")
             except Exception as e:
                 logger.error(f"Failed to generate scan.json from Google Sheets: {e}", exc_info=True)