Added execution manager and live display

prowler/lib/check/check.py

@@ -11,15 +11,15 @@ from types import ModuleType
 from typing import Any
 
 from colorama import Fore, Style
-from rich.live import Live
 
 import prowler
 from prowler.lib.check.compliance_models import load_compliance_framework
 from prowler.lib.check.custom_checks_metadata import update_check_metadata
+from prowler.lib.check.managers import ExecutionManager
 from prowler.lib.check.models import Check, load_check_metadata
 from prowler.lib.logger import logger
 from prowler.lib.outputs.outputs import report
-from prowler.lib.utils.ui import progress_manager
+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
@@ -492,44 +492,44 @@ def execute_checks(
         print(
             f"{Style.BRIGHT}Executing {checks_num} {check_noun}, please wait...{Style.RESET_ALL}\n"
         )
+        execution_manager = ExecutionManager(provider, checks_to_execute)
+        total_checks = execution_manager.total_checks_per_service()
+        completed_checks = {service: 0 for service in total_checks}
+        service_findings = []
+        for service, check_name in execution_manager.execute_checks():
+            try:
+                check_findings = execute(
+                    service,
+                    check_name,
+                    provider,
+                    audit_output_options,
+                    audit_info,
+                    services_executed,
+                    checks_executed,
+                    custom_checks_metadata,
+                )
+                all_findings.extend(check_findings)
+                service_findings.extend(check_findings)
+                # Update the completed checks count
+                completed_checks[service] += 1
 
-        check_dict = create_check_service_dict(checks_to_execute)
+                # Check if all checks for the service are completed
+                if completed_checks[service] == total_checks[service]:
+                    # All checks for the service are completed
+                    # Add a summary table or perform other actions
+                    live_display.add_summary_table_for_service(service_findings)
+                    # Clear service_findings
+                    service_findings = []
 
-        for service in check_dict:
-            service_check_num = len(check_dict[service])
-            progress_manager.create_manager(service, service_check_num)
-            current_manager = progress_manager.get_current_manager()
-            with Live(
-                current_manager.progress_table,
-                console=current_manager.console,
-                refresh_per_second=10,
-            ):
-                for check_name in check_dict[service]:
-                    try:
-                        check_findings = execute(
-                            service,
-                            check_name,
-                            provider,
-                            audit_output_options,
-                            audit_info,
-                            services_executed,
-                            checks_executed,
-                            custom_checks_metadata,
-                        )
-                        all_findings.extend(check_findings)
-
-                    # If check does not exists in the provider or is from another provider
-                    except ModuleNotFoundError:
-                        logger.error(
-                            f"Check '{check_name}' was not found for the {provider.upper()} provider"
-                        )
-                    except Exception as error:
-                        logger.error(
-                            f"{check_name} - {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
-                        )
-                    # Update the progress bar
-                    current_manager.clear_service_init_titles()
-                    current_manager.increment_progress()
+            # If check does not exists in the provider or is from another provider
+            except ModuleNotFoundError:
+                logger.error(
+                    f"Check '{check_name}' was not found for the {provider.upper()} provider"
+                )
+            except Exception as error:
+                logger.error(
+                    f"{check_name} - {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
+                )
     return all_findings
 
 

prowler/lib/check/managers.py

@@ -0,0 +1,126 @@
+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.utils.check_to_client_mapper import get_dependencies_for_checks
+
+
+class ExecutionManager:
+    def __init__(self, provider, checks_to_execute):
+        self.live_display = live_display
+        self.live_display.start()
+        self.provider = provider
+        self.loaded_clients = defaultdict(int)
+        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(
+            self.check_dependencies
+        )
+        self.services_queue = self.initialize_services_queue(self.check_dependencies)
+
+    def initialize_remaining_checks(self, check_dependencies):
+        remaining_checks = {}
+        for service, checks in check_dependencies.items():
+            for check_name, clients in checks.items():
+                remaining_checks[(service, check_name)] = clients
+        return remaining_checks
+
+    def initialize_services_queue(self, check_dependencies):
+        return list(check_dependencies.keys())
+
+    def total_checks_per_service(self):
+        """Returns a dictionary with the total number of checks for each service."""
+        total_checks = {}
+        for service, checks in self.check_dict.items():
+            total_checks[service] = len(checks)
+        return total_checks
+
+    def find_next_service(self):
+        # Prioritize services that use already loaded clients
+        for service in self.services_queue:
+            checks = self.check_dependencies[service]
+            if any(
+                client in self.loaded_clients and self.loaded_clients[client] > 0
+                for check in checks.values()
+                for client in check
+            ):
+                return service
+        return None if not self.services_queue else self.services_queue[0]
+
+    def import_client(self, client_name):
+        if self.loaded_clients[client_name] == 0:
+            # Dynamically import the client
+            module_name, _ = client_name.rsplit("_", 1)
+            client_module = importlib.import_module(
+                f"prowler.providers.{self.provider}.services.{module_name}.{client_name}"
+            )
+            setattr(self, client_name, client_module)
+        self.loaded_clients[client_name] += 1
+
+    def release_clients(self, completed_check_clients):
+        for client_name in completed_check_clients:
+            self.loaded_clients[client_name] -= 1
+            if self.loaded_clients[client_name] == 0 and not any(
+                client
+                for check in self.remaining_checks
+                for client in self.remaining_checks[check]
+            ):
+                del self.loaded_clients[client_name]
+                module_name, _ = client_name.rsplit("_", 1)
+                del sys.modules[
+                    f"prowler.providers.aws.services.{module_name}.{client_name}"
+                ]
+                # To check GC
+                weakref.finalize(getattr(self, client_name), on_finalize, client_name)
+                delattr(self, client_name)
+
+    def create_finalizer(self, client_name):
+        def on_finalize():
+            self.live_display.print_message(
+                f"Client {client_name} is being garbage collected."
+            )
+            print("gc")
+
+        return on_finalize
+
+    def execute_checks(self):
+        while self.remaining_checks:
+            next_service = self.find_next_service()
+            if not next_service:
+                break
+
+            if not self.live_display.has_section(next_service):
+                total_checks = len(self.check_dict[next_service])
+                self.live_display.add_service_section(next_service, total_checks)
+
+            self.services_queue.remove(next_service)
+            checks = self.check_dependencies[next_service]
+            for check_name in checks:
+                clients = checks[check_name]
+                for client in clients:
+                    self.import_client(client)
+
+                yield next_service, check_name
+
+                self.live_display.increment_check_progress()
+
+                self.release_clients(clients)
+                del self.remaining_checks[(next_service, check_name)]
+
+    @staticmethod
+    def create_check_service_dict(checks_to_execute):
+        output = {}
+        for check_name in checks_to_execute:
+            service = check_name.split("_")[0]
+            if service not in output.keys():
+                output[service] = []
+            output[service].append(check_name)
+        return output
+
+
+def on_finalize(client_name):
+    print(f"Client {client_name} is being garbage collected.")

prowler/lib/check/models.py

@@ -7,7 +7,7 @@ from pydantic import BaseModel, ValidationError
 from rich.progress import Task
 
 from prowler.lib.logger import logger
-from prowler.lib.utils.ui import progress_manager
+from prowler.lib.ui.live_display import live_display
 
 
 class Code(BaseModel):
@@ -79,28 +79,28 @@ class Check(ABC, Check_Metadata_Model):
         # Calls parents init function
         super().__init__(**data)
 
-        current_manager = progress_manager.get_current_manager()
+        current_section = live_display.get_current_section()
         # Cant do this as it messes with self.metdata()
-        # self.title_bar = current_manager.title_bar
-        # self.task_progress = current_manager.task_progress
+        # self.title_bar = current_section.title_bar
+        # self.task_progress = current_section.task_progress
 
-        self.title_bar_task = current_manager.title_bar.add_task(
+        self.title_bar_task = current_section.title_bar.add_task(
             f"{self.CheckTitle}...", start=False
         )
 
     def increment_task_progress(self):
-        current_manager = progress_manager.get_current_manager()
-        current_manager.task_progress.update(self.progress_task, advance=1)
+        current_section = live_display.get_current_section()
+        current_section.task_progress.update(self.progress_task, advance=1)
 
     def start_task(self, message, count):
-        current_manager = progress_manager.get_current_manager()
-        self.progress_task = current_manager.task_progress.add_task(
-            task_id=self.progress_task, description=message, total=count, visible=True
+        current_section = live_display.get_current_section()
+        self.progress_task = current_section.task_progress.add_task(
+            description=message, total=count, visible=True
         )
 
     def update_title_with_findings(self, findings):
-        current_manager = progress_manager.get_current_manager()
-        current_manager.task_progress.remove_task(self.progress_task)
+        current_section = live_display.get_current_section()
+        current_section.task_progress.remove_task(self.progress_task)
         total_failed = len([report for report in findings if report.status == "FAIL"])
         total_checked = len(findings)
         if total_failed == 0:
@@ -111,7 +111,7 @@ class Check(ABC, Check_Metadata_Model):
             message = (
                 f"{self.CheckTitle} [fail]{total_failed}/{total_checked} failed![/fail]"
             )
-        current_manager.title_bar.update(
+        current_section.title_bar.update(
             task_id=self.title_bar_task, description=message
         )
 

prowler/lib/ui/live_display.py

@@ -0,0 +1,190 @@
+import os
+import pathlib
+
+from rich.align import Align
+from rich.console import Console, Group
+from rich.live import Live
+from rich.panel import Panel
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+from rich.rule import Rule
+from rich.table import Table
+from rich.theme import Theme
+
+from prowler.lib.logger import logger
+
+
+class LiveDisplay(Live):
+    def __init__(self, *args, **kwargs):
+        theme = self.__load_theme_from_file__()
+        super().__init__(renderable=None, console=Console(theme=theme), *args, **kwargs)
+        self.sections = {}
+        self.ordered_section_names = []  # List to remember the order of sections
+        self.current_section = None
+
+    def has_section(self, section_name):
+        return section_name in self.sections.keys()
+
+    def add_service_section(self, service_name, total_checks):
+        # Create a new section for the service
+        if service_name in self.sections:
+            logger.error(f"Section already exists for {service_name}")
+            return
+        service_section = ServiceSection(service_name, total_checks)
+        self.sections[service_name] = service_section
+        self.ordered_section_names.append(service_name)  # Store the order
+        self.current_section = service_name
+        self.__update_layout__()
+
+    def get_current_section(self):
+        if not self.current_section:
+            logger.error(
+                "LiveDisplay has not been intialized! No current section to return"
+            )
+            return
+        return self.sections[self.current_section]
+
+    def print_message(self, message):
+        self.console.print(message)
+
+    def __update_layout__(self):
+        # Just a basic layout for now. Will be improved
+        # Create a group of renderables based on the order of sections
+        renderables = [
+            self.sections[name].renderables for name in self.ordered_section_names
+        ]
+        grouped_renderables = Group(*renderables)
+
+        # Update the existing layout
+        self.update(grouped_renderables)
+
+    def __load_theme_from_file__(self):
+        # Loads theme.yaml from the same folder as this file
+        actual_directory = pathlib.Path(os.path.dirname(os.path.realpath(__file__)))
+        with open(f"{actual_directory}/theme.yaml") as f:
+            theme = Theme.from_file(f)
+        return theme
+
+    # Wrappers to call the increment_progress in the ServiceSection class objects
+    def increment_check_progress(self):
+        current_section = self.get_current_section()
+        if not isinstance(current_section, ServiceSection):
+            logger.error("Current section is not set or not a ServiceSection")
+            return
+        current_section.increment_check_progress()
+
+    def add_summary_table_for_service(self, service_findings):
+        current_section = self.get_current_section()
+        if not isinstance(current_section, ServiceSection):
+            logger.error("Current section is not set or not a ServiceSection")
+            return
+        current_section.add_summary_table_for_service(service_findings)
+        self.update()
+
+
+class ServiceSection:
+    def __init__(self, service_name, total_checks) -> None:
+        self.service_name = service_name
+        self.total_checks = total_checks
+        self.renderables = self.__create_service_section__()
+        self.__start_check_progress__()
+
+    def __create_service_section__(self):
+        # Create the progress components
+        self.check_progress = Progress(
+            TextColumn("[bold]{task.description}"),
+            BarColumn(bar_width=None),
+            MofNCompleteColumn(),
+            transient=False,  # Optional: set True if you want the progress bar to disappear after completion
+        )
+
+        # Used to add titles that dont need progress bars
+        self.title_bar = Progress(
+            TextColumn("[progress.description]{task.description}"), transient=True
+        )
+        # Progress Bar for Service Init and Checks
+        self.task_progress = Progress(
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(bar_width=None),
+            MofNCompleteColumn(),
+            TimeElapsedColumn(),
+            TimeRemainingColumn(),
+            transient=True,
+        )
+
+        return Group(
+            Panel(
+                Group(
+                    self.check_progress,
+                    Rule(style="bold blue"),
+                    self.title_bar,
+                    Rule(style="bold blue"),
+                    self.task_progress,
+                ),
+                title=f"Service: {self.service_name}",
+            ),
+        )
+
+    def __start_check_progress__(self):
+        self.check_progress_task_id = self.check_progress.add_task(
+            "Checks executed", total=self.total_checks
+        )
+
+    def increment_check_progress(self):
+        self.check_progress.update(self.check_progress_task_id, advance=1)
+
+    def add_summary_table_for_service(self, service_findings):
+        # Calculate the total number of checks
+        total_checks = len(service_findings)
+
+        # Create a new renderable to show the results
+        results_table = Table(title="Service Check Results")
+        results_table.add_column("Result", justify="right")
+        results_table.add_column("Count", justify="center")
+        results_table.add_column("Percentage", justify="center")
+
+        # For each status type, determine the count and percentage
+        statuses = list(set([report.status for report in service_findings]))
+        for status in statuses:
+            count = len(
+                [report for report in service_findings if report.status == status]
+            )
+            percentage = (count / total_checks * 100) if total_checks else 0
+            results_table.add_row(
+                f"{status.capitalize()}",
+                f"[{status.lower()}]{str(count)}[/{status.lower()}]",
+                f"{percentage:.2f}%",
+            )
+
+        # Create a centered Panel
+        centered_results_table = Panel(Align.center(results_table))
+
+        # Replace the task_progress in the progress_table
+        self.renderables = Group(
+            Panel(
+                Group(
+                    self.check_progress,
+                    Rule(style="bold blue"),
+                    self.title_bar,
+                    Rule(style="bold blue"),
+                    centered_results_table,  # Replacing task_progress with results_table
+                ),
+                title=f"Service: {self.service_name}",
+            ),
+        )
+
+
+class ServiceInitSection:
+    def __init__(self) -> None:
+        pass
+
+
+# Create an instance of LiveDisplay to import elsewhere (ExecutionManager, the checks, the services)
+
+live_display = LiveDisplay()

prowler/lib/utils/check_to_client_mapper.py

@@ -0,0 +1,58 @@
+import ast
+import os
+import pathlib
+
+from prowler.lib.logger import logger
+
+
+class CheckFileFinder(ast.NodeVisitor):
+    def __init__(self):
+        self.is_check_file = False
+
+    def visit_ClassDef(self, node):
+        for base in node.bases:
+            if isinstance(base, ast.Name) and base.id == "Check":
+                self.is_check_file = True
+                break
+        self.generic_visit(node)
+
+
+class ImportFinder(ast.NodeVisitor):
+    def __init__(self, provider):
+        self.imports = set()
+        self.provider = provider
+
+    def visit_ImportFrom(self, node):
+        if node.module and f"prowler.providers.{self.provider}.services" in node.module:
+            for name in node.names:
+                if "_client" in name.name:
+                    self.imports.add(name.name)
+        self.generic_visit(node)
+
+
+def get_dependencies_for_checks(provider, checks_dict):
+    def analyze_check_file(file_path, provider):
+        # Prase the check file
+        with open(file_path, "r") as file:
+            node = ast.parse(file.read(), filename=file_path)
+
+        finder = ImportFinder(provider)
+        finder.visit(node)
+        return list(finder.imports)
+
+    current_directory = pathlib.Path(os.path.dirname(os.path.realpath(__file__)))
+    prowler_dir = current_directory.parent.parent
+    check_dependencies = {}
+    for service_name, checks in checks_dict.items():
+        check_dependencies[service_name] = {}
+        for check_name in checks:
+            relative_path = f"providers/{provider}/services/{service_name}/{check_name}/{check_name}.py"
+            check_file_path = prowler_dir / relative_path
+            if not check_file_path.exists():
+                logger.error(
+                    f"{check_name} does not exist at {relative_path}! Cannot determine service dependencies"
+                )
+                continue
+            clients = analyze_check_file(str(check_file_path), provider)
+            check_dependencies[service_name][check_name] = clients
+    return check_dependencies

prowler/lib/utils/ui.py

@@ -1,120 +0,0 @@
-import os
-import pathlib
-
-from rich.console import Console, Group
-from rich.panel import Panel
-from rich.progress import (
-    BarColumn,
-    MofNCompleteColumn,
-    Progress,
-    TextColumn,
-    TimeElapsedColumn,
-    TimeRemainingColumn,
-)
-from rich.rule import Rule
-from rich.theme import Theme
-
-
-class ProgressManager:
-    """
-    Keeps the progress_table for each service
-        This is passed to Live() to be rendered
-    overall_progress: tracks the checks progress
-    title_bar: used to let the user know what is happening/has happened
-    task_progress: shows a progress bar of what is happening
-    """
-
-    def __init__(self, service_name: str, total_checks: int, console: Console):
-        self.service_name = service_name
-        self.total_checks = total_checks
-        self.console = console
-        self.create_progress_table()
-
-    def create_progress_table(self):
-        # Create the progress components
-        self.overall_progress = Progress(
-            TextColumn("[bold]{task.description}"),
-            BarColumn(bar_width=None),
-            MofNCompleteColumn(),
-            transient=False,  # Optional: set True if you want the progress bar to disappear after completion
-        )
-
-        # Used to add titles that dont need progress bars
-        self.title_bar = Progress(
-            TextColumn("[progress.description]{task.description}"), transient=True
-        )
-        # Progress Bar for Service Init and Checks
-        self.task_progress = Progress(
-            TextColumn("[progress.description]{task.description}"),
-            BarColumn(bar_width=None),
-            MofNCompleteColumn(),
-            TimeElapsedColumn(),
-            TimeRemainingColumn(),
-            transient=True,
-        )
-
-        self.progress_table = Group(
-            Panel(
-                Group(
-                    self.overall_progress,
-                    Rule(style="bold blue"),
-                    self.title_bar,
-                    Rule(style="bold blue"),
-                    self.task_progress,
-                ),
-                title=f"Service: {self.service_name}",
-            ),
-        )
-
-        self.overall_progress_task_id = self.overall_progress.add_task(
-            "Checks executed", total=self.total_checks
-        )
-
-    def clear_service_init_titles(self):
-        service_init_tasks = [
-            task.id
-            for task in self.title_bar.tasks
-            if task.fields.get("task_type") == "Service"
-        ]
-        for task_id in service_init_tasks:
-            self.title_bar.remove_task(task_id)
-
-    def increment_progress(self):
-        self.overall_progress.update(self.overall_progress_task_id, advance=1)
-
-
-class ActiveProgressManager:
-    """
-    Used to keep track of the current progress manager, so that the current progress manager can be dynamically imported into services/checks
-    Also handles loading the theme, passing it to the console, and passing that console to the progress managers as they are created
-    Needs to be intialized before being imported into checks/services
-    """
-
-    def __init__(self):
-        theme = self.load_theme_from_file()
-        self.console = Console(theme=theme)
-        self.managers = {}
-        self.current_manager = None
-
-    def load_theme_from_file(self):
-        # Loads theme.yaml from the same folder as this file
-        actual_directory = pathlib.Path(os.path.dirname(os.path.realpath(__file__)))
-        with open(f"{actual_directory}/theme.yaml") as f:
-            theme = Theme.from_file(f)
-        return theme
-
-    def get_current_manager(self):
-        return self.current_manager
-
-    def set_current_manager(self, progress_manager: ProgressManager):
-        self.current_manager = progress_manager
-
-    def create_manager(self, service_name: str, total_checks: int):
-        if service_name not in self.managers:
-            self.managers[service_name] = ProgressManager(
-                service_name, total_checks, self.console
-            )
-        self.current_manager = self.managers[service_name]
-
-
-progress_manager = ActiveProgressManager()

prowler/providers/aws/lib/service/service.py

@@ -1,7 +1,7 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
 from prowler.lib.logger import logger
-from prowler.lib.utils.ui import progress_manager
+from prowler.lib.ui.live_display import live_display
 from prowler.providers.aws.aws_provider import (
     generate_regional_clients,
     get_default_region,
@@ -51,10 +51,10 @@ class AWSService:
         self.thread_pool = ThreadPoolExecutor(max_workers=MAX_WORKERS)
 
         # Progress bar to add tasks to
-        current_manager = progress_manager.get_current_manager()
-        self.progress = current_manager.task_progress
+        current_section = live_display.get_current_section()
+        self.progress = current_section.task_progress
         self.progress_tasks = []
-        self.title_bar = current_manager.title_bar
+        self.title_bar = current_section.title_bar
 
         self.title_bar_task = self.title_bar.add_task(
             f"Intializing {self.service} service:", start=False, task_type="Service"