ben.carrasco/prowler
1
0
Fork 0
mirror of https://github.com/prowler-cloud/prowler.git synced 2026-07-24 13:01:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

cleaned up execution manager,live display. Added metaclass

Browse Source
This commit is contained in:
Fennerr
2023-12-20 23:22:18 +02:00
parent f2f922d7e8
commit 0d97780ade
4 changed files with 217 additions and 272 deletions
Download Patch File Download Diff File Expand all files Collapse all files
prowler/lib/check/managers.py
+133 -158
View File
@@ -3,7 +3,7 @@ import os
import sys
import traceback
from types import ModuleType
from typing import Any
from typing import Any, Set
from colorama import Fore, Style
@@ -43,8 +43,21 @@ class ExecutionManager:
)
self.services_queue = self.initialize_services_queue(self.check_dependencies)
self.services_executed = (set(),)
self.checks_executed = (set(),)
# For tracking the executed services and checks
self.services_executed: Set[str] = set()
self.checks_executed: Set[str] = set()
# Initialize the Audit Metadata
self.audit_info.audit_metadata = Audit_Metadata(
services_scanned=0,
expected_checks=self.checks_to_execute,
completed_checks=0,
audit_progress=0,
)
def update_tracking(self, service: str, check: str):
self.services_executed.add(service)
self.checks_executed.add(check)
@staticmethod
def initialize_remaining_checks(check_dependencies):
@@ -58,6 +71,16 @@ class ExecutionManager:
def initialize_services_queue(check_dependencies):
return list(check_dependencies.keys())
@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 total_checks_per_service(self):
"""Returns a dictionary with the total number of checks for each service."""
total_checks = {}
@@ -77,18 +100,6 @@ class ExecutionManager:
return service
return None if not self.services_queue else self.services_queue[0]
# Imports service clients, and tracks if it needs to be imported
def import_client(self, client_name):
if not self.loaded_clients.get(client_name):
self.live_display.add_client_init_section(client_name)
# 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}"
)
self.loaded_clients[client_name] = client_module
self.live_display.remove_client_init_section()
@staticmethod
def import_check(check_path: str) -> ModuleType:
"""
@@ -102,6 +113,16 @@ class ExecutionManager:
lib = importlib.import_module(f"{check_path}")
return lib
# Imports service clients, and tracks if it needs to be imported
def import_client(self, client_name):
if not self.loaded_clients.get(client_name):
# 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}"
)
self.loaded_clients[client_name] = client_module
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
@@ -144,11 +165,10 @@ class ExecutionManager:
for client in clients_for_service:
self.live_display.add_client_init_section(client)
self.import_client(client)
self.live_display.remove_client_init_section()
if not self.live_display.has_section(current_service):
total_checks = len(self.check_dict[current_service])
self.live_display.add_service_section(current_service, total_checks)
# Add the display component
total_checks = len(self.check_dict[current_service])
self.live_display.add_service_section(current_service, total_checks)
for check_name, clients_for_check in checks.items():
@@ -161,24 +181,87 @@ class ExecutionManager:
self.release_clients(clients_for_check)
self.live_display.increment_overall_service_progress()
self.live_display.remove_section(current_service)
@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 execute_checks(self) -> list:
# List to store all the check's findings
all_findings = []
if os.name != "nt":
try:
from resource import RLIMIT_NOFILE, getrlimit
# Check ulimit for the maximum system open files
soft, _ = getrlimit(RLIMIT_NOFILE)
if soft < 4096:
logger.warning(
f"Your session file descriptors limit ({soft} open files) is below 4096. We recommend to increase it to avoid errors. Solve it running this command `ulimit -n 4096`. For more info visit https://docs.prowler.cloud/en/latest/troubleshooting/"
)
except Exception as error:
logger.error("Unable to retrieve ulimit default settings")
logger.error(
f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
)
# Execution with the --only-logs flag
if self.audit_output_options.only_logs:
for service, check_name in self.generate_checks():
try:
check_findings = self.execute(service, check_name)
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 {self.provider.upper()} provider"
)
except Exception as error:
logger.error(
f"{check_name} - {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
)
else:
# Default execution
total_checks = self.total_checks_per_service()
self.live_display.add_overall_progress_section(
total_checks_dict=total_checks
)
# For tracking when a service is completed
completed_checks = {service: 0 for service in total_checks}
service_findings = []
for service, check_name in self.generate_checks():
try:
check_findings = self.execute(
service,
check_name,
)
all_findings.extend(check_findings)
service_findings.extend(check_findings)
# Update the completed checks count
completed_checks[service] += 1
# 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_results_for_service(service, service_findings)
# Clear service_findings
service_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 {self.provider.upper()} provider"
)
except Exception as error:
logger.error(
f"{check_name} - {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
)
self.live_display.hide_service_section()
return all_findings
def execute(
self,
service: str,
check_name: str,
services_executed: set,
checks_executed: set,
):
# Import check module
check_module_path = f"prowler.providers.{self.provider}.services.{service}.{check_name}.{check_name}"
@@ -199,11 +282,8 @@ class ExecutionManager:
check_findings = self.run_check(c, self.audit_output_options)
# Update Audit Status
services_executed.add(service)
checks_executed.add(check_name)
self.audit_info.audit_metadata = self.update_audit_metadata(
self.audit_info.audit_metadata, services_executed, checks_executed
)
self.update_tracking(service, check_name)
self.update_audit_metadata()
# Allowlist findings
if self.audit_output_options.allowlist_file:
@@ -231,28 +311,6 @@ class ExecutionManager:
return check_findings
@staticmethod
def update_audit_metadata(
audit_metadata: Audit_Metadata, services_executed: set, checks_executed: set
) -> Audit_Metadata:
"""update_audit_metadata returns the audit_metadata updated with the new status
Updates the given audit_metadata using the length of the services_executed and checks_executed
"""
try:
audit_metadata.services_scanned = len(services_executed)
audit_metadata.completed_checks = len(checks_executed)
audit_metadata.audit_progress = (
100 * len(checks_executed) / len(audit_metadata.expected_checks)
)
return audit_metadata
except Exception as error:
logger.error(
f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
)
@staticmethod
def run_check(check: Check, output_options: Provider_Output_Options) -> list:
findings = []
@@ -274,105 +332,22 @@ class ExecutionManager:
finally:
return findings
def execute_checks(self) -> list:
# List to store all the check's findings
all_findings = []
# Services and checks executed for the Audit Status
services_executed = set()
checks_executed = set()
def update_audit_metadata(self):
"""update_audit_metadata returns the audit_metadata updated with the new status
# Initialize the Audit Metadata
self.audit_info.audit_metadata = Audit_Metadata(
services_scanned=0,
expected_checks=self.checks_to_execute,
completed_checks=0,
audit_progress=0,
)
if os.name != "nt":
try:
from resource import RLIMIT_NOFILE, getrlimit
# Check ulimit for the maximum system open files
soft, _ = getrlimit(RLIMIT_NOFILE)
if soft < 4096:
logger.warning(
f"Your session file descriptors limit ({soft} open files) is below 4096. We recommend to increase it to avoid errors. Solve it running this command `ulimit -n 4096`. For more info visit https://docs.prowler.cloud/en/latest/troubleshooting/"
)
except Exception as error:
logger.error("Unable to retrieve ulimit default settings")
logger.error(
f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
)
# Execution with the --only-logs flag
if self.audit_output_options.only_logs:
for check_name in self.checks_to_execute:
# Recover service from check name
service = check_name.split("_")[0]
try:
check_findings = self.execute(
service,
check_name,
self.provider,
self.audit_output_options,
self.audit_info,
services_executed,
checks_executed,
self.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 {self.provider.upper()} provider"
)
except Exception as error:
logger.error(
f"{check_name} - {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
)
else:
# Default execution
total_checks = self.total_checks_per_service()
self.live_display.add_overall_progress_section(
total_checks_dict=total_checks
Updates the given audit_metadata using the length of the services_executed and checks_executed
"""
try:
self.audit_info.audit_metadata.services_scanned = len(
self.services_executed
)
self.audit_info.audit_metadata.completed_checks = len(self.checks_executed)
self.audit_info.audit_metadata.audit_progress = (
100
* len(self.checks_executed)
/ len(self.audit_info.audit_metadata.expected_checks)
)
except Exception as error:
logger.error(
f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
)
# For tracking when a service is completed
completed_checks = {service: 0 for service in total_checks}
service_findings = []
for service, check_name in self.generate_checks():
try:
check_findings = self.execute(
service,
check_name,
services_executed,
checks_executed,
)
all_findings.extend(check_findings)
service_findings.extend(check_findings)
# Update the completed checks count
completed_checks[service] += 1
# 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_results_for_service(service, service_findings)
# Clear service_findings
service_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 {self.provider.upper()} provider"
)
except Exception as error:
logger.error(
f"{check_name} - {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
)
return all_findings
def on_finalize(client_name):
print(f"Client {client_name} is being garbage collected.")
prowler/lib/check/models.py
+43 -9
View File
@@ -2,8 +2,10 @@ import os
import sys
from abc import ABC, abstractmethod
from dataclasses import dataclass
from functools import wraps
from pydantic import BaseModel, ValidationError
from pydantic.main import ModelMetaclass
from rich.progress import Task
from prowler.lib.logger import logger
@@ -59,7 +61,24 @@ class Check_Metadata_Model(BaseModel):
Compliance: list = None
class Check(ABC, Check_Metadata_Model):
class CheckMeta(ModelMetaclass):
"""
Dynamically decorates the execute function of all subclasses of the Check class
By making CheckMeta inherit from ModelMetaclass, it ensures that all features provided by Pydantic's BaseModel (such as data validation, serialization, and so forth) are preserved. CheckMeta just adds additional behavior (decorator application) on top of the existing features.
This also works because ModelMetaclass inherits from ABCMeta, as does the ABC class (its got to do with how metaclasses work when applying it to a class that inherits from other classes that have a metaclass).
The primary role of CheckMeta is to automatically apply a decorator to the execute method of subclasses. This behavior does not conflict with the typical responsibilities of ModelMetaclass
"""
def __new__(cls, name, bases, dct):
if "execute" in dct and not getattr(
dct["execute"], "__isabstractmethod__", False
):
dct["execute"] = Check.update_title_with_findings_decorator(dct["execute"])
return super(CheckMeta, cls).__new__(cls, name, bases, dct)
class Check(ABC, Check_Metadata_Model, metaclass=CheckMeta):
"""Prowler Check"""
title_bar_task: Task = None
@@ -79,28 +98,25 @@ class Check(ABC, Check_Metadata_Model):
# Calls parents init function
super().__init__(**data)
current_section = live_display.get_current_section()
# Cant do this as it messes with self.metdata()
# self.title_bar = current_section.title_bar
# self.task_progress = current_section.task_progress
current_section = live_display.get_service_section()
self.title_bar_task = current_section.title_bar.add_task(
f"{self.CheckTitle}...", start=False
)
def increment_task_progress(self):
current_section = live_display.get_current_section()
current_section = live_display.get_service_section()
current_section.task_progress.update(self.progress_task, advance=1)
def start_task(self, message, count):
current_section = live_display.get_current_section()
current_section = live_display.get_service_section()
self.progress_task = current_section.task_progress.add_task(
description=message, total=count, visible=True
)
def update_title_with_findings(self, findings):
current_section = live_display.get_current_section()
current_section.task_progress.remove_task(self.progress_task)
current_section = live_display.get_service_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:
@@ -123,6 +139,24 @@ class Check(ABC, Check_Metadata_Model):
def execute(self):
"""Execute the check's logic"""
@staticmethod
def update_title_with_findings_decorator(func):
"""
Decorator to update the title bar in the live_display with findings after executing a check.
"""
@wraps(func)
def wrapper(check_instance, *args, **kwargs):
# Execute the original check's logic
findings = func(check_instance, *args, **kwargs)
# Update the title bar with the findings
check_instance.update_title_with_findings(findings)
return findings
return wrapper
@dataclass
class Check_Report:
prowler/lib/ui/live_display.py
+37 -104
View File
@@ -23,7 +23,6 @@ from rich.text import Text
from rich.theme import Theme
from prowler.config.config import prowler_version, timestamp
from prowler.lib.logger import logger
from prowler.providers.aws.lib.audit_info.models import AWS_Audit_Info
@@ -32,56 +31,37 @@ class LiveDisplay(Live):
theme = self.__load_theme_from_file__()
super().__init__(renderable=None, console=Console(theme=theme), *args, **kwargs)
self.sections = {}
self.ordered_service_names = [] # List to remember the order of sections
self.current_section = None
self.make_layout()
def has_section(self, section_name):
return section_name in self.sections.keys()
def add_section(self, section_name, section_class, default_section=False):
self.sections[section_name] = section_class
if not default_section:
self.ordered_service_names.append(section_name) # Store the order
self.current_section = section_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 get_section(self, section_name):
return self.sections[section_name]
def get_service_section(self):
# Used by Check
return self.sections["client_and_service"]
def get_client_init_section(self):
return self.sections["client_init"]
def remove_section(self, section_name):
del self.sections[section_name]
if section_name in self.ordered_service_names:
self.ordered_service_names.remove(section_name)
# Used by AWSService
return self.sections["client_and_service"]
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
# Used to create the ServiceSection when checks start to execute (after clients have been imported)
service_section = ServiceSection(service_name, total_checks)
self.add_section(service_name, service_section)
self.layout["client_init"].update(service_section)
self.sections["client_and_service"] = service_section
self.layout["client_and_service"].update(service_section)
def hide_service_section(self):
# To hide the last service after execution has completed
self.layout["client_and_service"].visible = False
def print_message(self, message):
# No use yet
self.console.print(message)
def make_layout(self):
"""
Define the layout.
Defines the layout.
Making sections invisible so it doesnt show the default Layout metadata before content is added
Text(" ") is to stop the layout metadata from rendering before the layout is updated with real content
client_and_service handles client init (when importing clients) and service check execution
"""
self.layout = Layout(name="root")
self.layout.split(
@@ -95,49 +75,11 @@ class LiveDisplay(Live):
)
self.layout["main"].split_row(
Layout(
Text(" "), name="client_init", ratio=3
Text(" "), name="client_and_service", ratio=3
), # For client_init and service
Layout(name="results", ratio=2, visible=False),
)
def __update_layout__(self):
for name, section in self.sections.items():
if self.layout.get(name):
if hasattr(section, "__rich__"):
self.layout[name].update(section)
self.layout[name].visible = True
else:
self.layout[name].visible = False
else:
# Its needs to be part of service section
pass
service_renderables = []
for name in self.ordered_service_names:
section_renderable = self.sections[name].renderables
if isinstance(section_renderable, list):
# If it's a list, create a Panel with these renderables
section_title = getattr(
self.sections[name], "title", None
) # Get title if it exists
panel = Panel(Group(*section_renderable), title=section_title)
service_renderables.append(panel)
else:
# Otherwise, treat it as a single renderable
service_renderables.append(section_renderable)
if service_renderables:
grouped_renderables = Group(*service_renderables)
self.layout["service"].visible = True
self.layout["service"].update(grouped_renderables)
else:
self.layout["service"].visible = False
# Update the existing layout
self.console.print(self.layout.tree)
self.update(self.layout)
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__)))
@@ -147,41 +89,32 @@ class LiveDisplay(Live):
# Wrappers for ServiceSection methods
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()
service_section = self.sections["client_and_service"]
service_section.increment_check_progress()
# Results Section Methods
def add_results_for_service(self, service_name, service_findings):
results_section = self.sections["results"]
results_section.add_results_for_service(service_name, service_findings)
def add_results_section(self):
# Intializes the results section
results_layout = self.layout["results"]
results_section = ResultsSection()
results_layout.update(results_section)
results_layout.visible = True
self.sections["results"] = results_section
def add_results_for_service(self, service_name, service_findings):
# Adds rows to the Service Check Results table
results_section = self.sections["results"]
results_section.add_results_for_service(service_name, service_findings)
# Client Init Methods
def add_client_init_section(self, service_name):
client_init_section = ClientInitSection(service_name)
self.add_section("client_init", client_init_section, default_section=True)
self.layout["client_init"].update(client_init_section)
self.layout["client_init"].visible = True
# self.add_section("client_and_service", client_init_section, default_section=True)
self.sections["client_and_service"] = client_init_section
self.layout["client_and_service"].update(client_init_section)
self.layout["client_and_service"].visible = True
self.update(self.layout)
def remove_client_init_section(self):
return
if "client_init" not in self.sections:
# Might not have needed to init any services
return
self.remove_section("client_init")
# self.layout["client_init"].visible = False
self.__update_layout__()
# Intro Section Methods
def add_intro(self, args):
# A way to get around parsing args to LiveDisplay when it is intialized
@@ -189,34 +122,34 @@ class LiveDisplay(Live):
self.cli_args = args
intro_layout = self.layout["intro"]
intro_section = IntroSection(args, intro_layout)
self.add_section("intro", intro_section, default_section=True)
self.sections["intro"] = intro_section
self.update(self.layout)
def print_aws_credentials(self, audit_info):
intro_section = self.get_section("intro")
# Adds the AWS credentials to the display - will need to extend to gcp and azure
intro_section = self.sections["intro"]
intro_section.add_aws_credentials(audit_info)
# Overall Progress Methods
def add_overall_progress_section(self, total_checks_dict):
# section = self.get_section("intro")
overall_progress_section = OverallProgressSection(total_checks_dict)
overall_progress_layout = self.layout["overall_progress"]
overall_progress_layout.update(overall_progress_section)
overall_progress_layout.visible = True
self.add_section(
"overall_progress", overall_progress_section, default_section=True
)
self.sections["overall_progress"] = overall_progress_section
# Add results section
self.add_results_section()
self.update(self.layout)
def increment_overall_check_progress(self):
section = self.get_section("overall_progress")
# Called by ExecutionManager
section = self.sections["overall_progress"]
section.increment_check_progress()
def increment_overall_service_progress(self):
section = self.get_section("overall_progress")
# Called by ExecutionManager
section = self.sections["overall_progress"]
section.increment_service_progress()
@@ -268,7 +201,7 @@ class ServiceSection:
def __start_check_progress__(self):
self.check_progress_task_id = self.check_progress.add_task(
f"{self.service_name} checks executed", total=self.total_checks
"Checks executed", total=self.total_checks
)
def increment_check_progress(self):
prowler/providers/aws/lib/service/service.py
+4 -1
View File
@@ -106,7 +106,10 @@ class AWSService:
@staticmethod
def progress_decorator(func):
"""Decorator to update the progress bar before and after a function call."""
"""
Decorator to update the progress bar before and after a function call.
To be used for methods within global services, which do not make use of the __threading_call__ function
"""
@wraps(func)
def wrapper(self, *args, **kwargs):