diff --git a/prowler/__main__.py b/prowler/__main__.py index 57aa643171..0d35c56f87 100644 --- a/prowler/__main__.py +++ b/prowler/__main__.py @@ -12,7 +12,6 @@ from prowler.lib.check.check import ( bulk_load_compliance_frameworks, exclude_checks_to_run, exclude_services_to_run, - execute_checks, list_categories, list_checks_json, list_services, @@ -30,6 +29,7 @@ from prowler.lib.check.custom_checks_metadata import ( parse_custom_checks_metadata_file, update_checks_metadata, ) +from prowler.lib.check.managers import ExecutionManager from prowler.lib.cli.parser import ProwlerArgumentParser from prowler.lib.logger import logger, set_logging_config from prowler.lib.outputs.compliance import display_compliance_table @@ -54,7 +54,6 @@ from prowler.providers.common.audit_info import ( from prowler.providers.common.outputs import set_provider_output_options from prowler.providers.common.quick_inventory import run_provider_quick_inventory -from prowler.lib.check.managers import ExecutionManager def prowler(): # Parse Arguments @@ -188,7 +187,7 @@ def prowler(): # Execute checks findings = [] - + if len(checks_to_execute): execution_manager = ExecutionManager( checks_to_execute, diff --git a/prowler/lib/check/check_to_client_mapper.py b/prowler/lib/check/check_to_client_mapper.py index 6fcfa8f294..5c2025c1db 100644 --- a/prowler/lib/check/check_to_client_mapper.py +++ b/prowler/lib/check/check_to_client_mapper.py @@ -4,6 +4,7 @@ import pathlib from prowler.lib.logger import logger + class ImportFinder(ast.NodeVisitor): def __init__(self, provider): self.imports = set() @@ -16,6 +17,7 @@ class ImportFinder(ast.NodeVisitor): self.imports.add(name.name) self.generic_visit(node) + def analyze_check_file(file_path, provider): # Prase the check file with open(file_path, "r") as file: @@ -25,6 +27,7 @@ def analyze_check_file(file_path, provider): finder.visit(node) return list(finder.imports) + def get_dependencies_for_checks(provider, checks_dict): current_directory = pathlib.Path(os.path.dirname(os.path.realpath(__file__))) diff --git a/prowler/lib/check/managers.py b/prowler/lib/check/managers.py index b62d45fc2d..ddaaa35b62 100644 --- a/prowler/lib/check/managers.py +++ b/prowler/lib/check/managers.py @@ -1,45 +1,44 @@ -import importlib -import sys - -# To check if client is being GC -import weakref -from collections import defaultdict - -from prowler.lib.ui.live_display import live_display -from prowler.lib.check.check_to_client_mapper import get_dependencies_for_checks - import importlib import os import sys import traceback -from pkgutil import walk_packages + +# To check if client is being GC +import weakref from types import ModuleType from typing import Any from colorama import Fore, Style +from prowler.lib.check.check_to_client_mapper import get_dependencies_for_checks from prowler.lib.check.custom_checks_metadata import update_check_metadata from prowler.lib.check.models import Check from prowler.lib.logger import logger from prowler.lib.outputs.outputs import report from prowler.lib.ui.live_display import live_display -from prowler.lib.utils.utils import open_file, parse_json_file from prowler.providers.aws.lib.allowlist.allowlist import allowlist_findings from prowler.providers.common.models import Audit_Metadata from prowler.providers.common.outputs import Provider_Output_Options class ExecutionManager: - def __init__(self, checks_to_execute: list, provider: str, audit_info: Any, audit_output_options: Provider_Output_Options, custom_checks_metadata: Any): + def __init__( + self, + checks_to_execute: list, + provider: str, + audit_info: Any, + audit_output_options: Provider_Output_Options, + custom_checks_metadata: Any, + ): self.checks_to_execute = checks_to_execute self.provider = provider self.audit_info = audit_info self.audit_output_options = audit_output_options self.custom_checks_metadata = custom_checks_metadata - + self.live_display = live_display self.live_display.start() - self.loaded_clients = {} #defaultdict(lambda: False) + self.loaded_clients = {} # defaultdict(lambda: False) self.check_dict = self.create_check_service_dict(checks_to_execute) self.check_dependencies = get_dependencies_for_checks(provider, self.check_dict) self.remaining_checks = self.initialize_remaining_checks( @@ -47,8 +46,8 @@ class ExecutionManager: ) self.services_queue = self.initialize_services_queue(self.check_dependencies) - self.services_executed = set(), - self.checks_executed = set(), + self.services_executed = (set(),) + self.checks_executed = (set(),) @staticmethod def initialize_remaining_checks(check_dependencies): @@ -90,8 +89,7 @@ class ExecutionManager: f"prowler.providers.{self.provider}.services.{module_name}.{client_name}" ) self.loaded_clients[client_name] = client_module - - + @staticmethod def import_check(check_path: str) -> ModuleType: """ @@ -107,14 +105,16 @@ class ExecutionManager: def release_clients(self, completed_check_clients): for client_name in completed_check_clients: - # Determine if any of the remaining checks still require the client + # Determine if any of the remaining checks still require the client if not any( client == client_name for check in self.remaining_checks for client in self.remaining_checks[check] ): # DEBUG: To check GC - weakref.finalize(self.loaded_clients[client_name], on_finalize, client_name) + weakref.finalize( + self.loaded_clients[client_name], on_finalize, client_name + ) # Delete the reference to the client for this object del self.loaded_clients[client_name] module_name, _ = client_name.rsplit("_", 1) @@ -133,7 +133,6 @@ class ExecutionManager: return on_finalize - def generate_checks(self): """ This is a generator function, which will: @@ -142,7 +141,7 @@ class ExecutionManager: * Yield the service and check name, 1-by-1, to be used within execute_checks * Pass the completed checks to release_clients to determine if the clients that were required by the check are no longer needed, and can be garabage collected It will complete the checks for a service, before moving onto the next one - It uses find_next_service to prioritize the next service based on if any of that service's checks require a client that has already been loaded + It uses find_next_service to prioritize the next service based on if any of that service's checks require a client that has already been loaded """ while self.remaining_checks: current_service = self.find_next_service() @@ -168,7 +167,6 @@ class ExecutionManager: del self.remaining_checks[(current_service, check_name)] self.release_clients(clients) - @staticmethod def create_check_service_dict(checks_to_execute): @@ -179,7 +177,7 @@ class ExecutionManager: output[service] = [] output[service].append(check_name) return output - + def execute( self, service: str, @@ -188,17 +186,19 @@ class ExecutionManager: checks_executed: set, ): # Import check module - check_module_path = ( - f"prowler.providers.{self.provider}.services.{service}.{check_name}.{check_name}" - ) + check_module_path = f"prowler.providers.{self.provider}.services.{service}.{check_name}.{check_name}" lib = self.import_check(check_module_path) # Recover functions from check check_to_execute = getattr(lib, check_name) c = check_to_execute() # Update check metadata to reflect that in the outputs - if self.custom_checks_metadata and self.custom_checks_metadata["Checks"].get(c.CheckID): - c = update_check_metadata(c, self.custom_checks_metadata["Checks"][c.CheckID]) + if self.custom_checks_metadata and self.custom_checks_metadata["Checks"].get( + c.CheckID + ): + c = update_check_metadata( + c, self.custom_checks_metadata["Checks"][c.CheckID] + ) # Run check check_findings = self.run_check(c, self.audit_output_options) @@ -228,12 +228,14 @@ class ExecutionManager: outputs_module = importlib.import_module(lib) custom_report_interface = getattr(outputs_module, "report") - custom_report_interface(check_findings, self.audit_output_options, self.audit_info) + custom_report_interface( + check_findings, self.audit_output_options, self.audit_info + ) except Exception: sys.exit(1) return check_findings - + @staticmethod def update_audit_metadata( audit_metadata: Audit_Metadata, services_executed: set, checks_executed: set @@ -276,10 +278,8 @@ class ExecutionManager: ) finally: return findings - - def execute_checks( - self - ) -> list: + + def execute_checks(self) -> list: # List to store all the check's findings all_findings = [] # Services and checks executed for the Audit Status @@ -382,5 +382,6 @@ class ExecutionManager: ) return all_findings + def on_finalize(client_name): - print(f"Client {client_name} is being garbage collected.") \ No newline at end of file + print(f"Client {client_name} is being garbage collected.")