From 4da6d152c3474b5beb7782ba76579a1d220c6146 Mon Sep 17 00:00:00 2001 From: Sergio Garcia <38561120+sergargar@users.noreply.github.com> Date: Thu, 13 Apr 2023 13:44:25 +0200 Subject: [PATCH] feat(custom checks): add -x/--checks-folder for custom checks (#2191) --- docs/tutorials/misc.md | 14 ++++ prowler/__main__.py | 17 ++++- prowler/lib/check/check.py | 63 ++++++++++++++++ prowler/lib/cli/parser.py | 6 ++ prowler/lib/outputs/html.py | 56 +++++++------- prowler/lib/outputs/models.py | 35 +++++---- tests/lib/check/check_test.py | 74 +++++++++++++++++++ .../checks_folder/check11/__init__.py | 0 .../checks_folder/check12/__init__.py | 0 .../checks_folder/check7777/__init__.py | 0 tests/lib/cli/parser_test.py | 17 +++++ tests/lib/outputs/outputs_test.py | 4 + 12 files changed, 244 insertions(+), 42 deletions(-) create mode 100644 tests/lib/check/fixtures/checks_folder/check11/__init__.py create mode 100644 tests/lib/check/fixtures/checks_folder/check12/__init__.py create mode 100644 tests/lib/check/fixtures/checks_folder/check7777/__init__.py diff --git a/docs/tutorials/misc.md b/docs/tutorials/misc.md index 8c9ed619..b75d3ae2 100644 --- a/docs/tutorials/misc.md +++ b/docs/tutorials/misc.md @@ -51,7 +51,21 @@ prowler -e/--excluded-checks ec2 rds ```console prowler -C/--checks-file .json ``` +## Custom Checks +Prowler allows you to include your custom checks with the flag: +```console +prowler -x/--checks-folder +``` +> S3 URIs are also supported as folders for custom checks, e.g. s3://bucket/prefix/checks_folder/. Make sure that the used credentials have s3:GetObject permissions in the S3 path where the custom checks are located. +The custom checks folder must contain one subfolder per check, each subfolder must be named as the check and must contain: + +- An empty `__init__.py`: to make Python treat this check folder as a package. +- A `check_name.py` containing the check's logic. +- A `check_name.metadata.json` containing the check's metadata. +>The check name must start with the service name followed by an underscore (e.g., ec2_instance_public_ip). + +To see more information about how to write checks see the [Developer Guide](../developer-guide/#create-a-new-check-for-a-provider). ## Severities Each of Prowler's checks has a severity, which can be: - informational diff --git a/prowler/__main__.py b/prowler/__main__.py index ca5dc3b9..15e0c185 100644 --- a/prowler/__main__.py +++ b/prowler/__main__.py @@ -12,11 +12,13 @@ from prowler.lib.check.check import ( execute_checks, list_categories, list_services, + parse_checks_from_folder, print_categories, print_checks, print_compliance_frameworks, print_compliance_requirements, print_services, + remove_custom_checks_module, ) from prowler.lib.check.checks_loader import load_checks_to_execute from prowler.lib.check.compliance import update_checks_metadata_with_compliance @@ -52,9 +54,17 @@ def prowler(): services = args.services categories = args.categories checks_file = args.checks_file + checks_folder = args.checks_folder severities = args.severity compliance_framework = args.compliance + # Set the audit info based on the selected provider + audit_info = set_provider_audit_info(provider, args.__dict__) + + # Import custom checks from folder + if checks_folder: + parse_checks_from_folder(audit_info, checks_folder, provider) + # We treat the compliance framework as another output format if compliance_framework: args.output_modes.extend(compliance_framework) @@ -126,9 +136,6 @@ def prowler(): print_checks(provider, checks_to_execute, bulk_checks_metadata) sys.exit() - # Set the audit info based on the selected provider - audit_info = set_provider_audit_info(provider, args.__dict__) - # Once the audit_info is set and we have the eventual checks based on the resource identifier, # it is time to check what Prowler's checks are going to be executed if audit_info.audit_resources: @@ -217,6 +224,10 @@ def prowler(): audit_output_options.output_directory, ) + # If custom checks were passed, remove the modules + if checks_folder: + remove_custom_checks_module(checks_folder, provider) + # If there are failed findings exit code 3, except if -z is input if not args.ignore_exit_code_3 and stats["total_fail"] > 0: sys.exit(3) diff --git a/prowler/lib/check/check.py b/prowler/lib/check/check.py index 893ce0d2..4214803e 100644 --- a/prowler/lib/check/check.py +++ b/prowler/lib/check/check.py @@ -1,6 +1,8 @@ import functools import importlib import os +import re +import shutil import sys import traceback from pkgutil import walk_packages @@ -24,6 +26,7 @@ except KeyError: except Exception: sys.exit(1) +import prowler from prowler.lib.utils.utils import open_file, parse_json_file from prowler.providers.common.models import Audit_Metadata from prowler.providers.common.outputs import Provider_Output_Options @@ -117,6 +120,66 @@ def parse_checks_from_file(input_file: str, provider: str) -> set: return checks_to_execute +# Load checks from custom folder +def parse_checks_from_folder(audit_info, input_folder: str, provider: str) -> int: + try: + imported_checks = 0 + # Check if input folder is a S3 URI + if provider == "aws" and re.search( + "^s3://([^/]+)/(.*?([^/]+))/$", input_folder + ): + bucket = input_folder.split("/")[2] + key = ("/").join(input_folder.split("/")[3:]) + s3_reource = audit_info.audit_session.resource("s3") + bucket = s3_reource.Bucket(bucket) + for obj in bucket.objects.filter(Prefix=key): + if not os.path.exists(os.path.dirname(obj.key)): + os.makedirs(os.path.dirname(obj.key)) + bucket.download_file(obj.key, obj.key) + input_folder = key + # Import custom checks by moving the checks folders to the corresponding services + with os.scandir(input_folder) as checks: + for check in checks: + if check.is_dir(): + check_module = input_folder + "/" + check.name + # Copy checks to specific provider/service folder + check_service = check.name.split("_")[0] + prowler_dir = prowler.__path__ + prowler_module = f"{prowler_dir[0]}/providers/{provider}/services/{check_service}/{check.name}" + if os.path.exists(prowler_module): + shutil.rmtree(prowler_module) + shutil.copytree(check_module, prowler_module) + imported_checks += 1 + return imported_checks + except Exception as error: + logger.critical( + f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}] -- {error}" + ) + sys.exit(1) + + +# Load checks from custom folder +def remove_custom_checks_module(input_folder: str, provider: str): + # Check if input folder is a S3 URI + s3_uri = False + if provider == "aws" and re.search("^s3://([^/]+)/(.*?([^/]+))/$", input_folder): + input_folder = ("/").join(input_folder.split("/")[3:]) + s3_uri = True + + with os.scandir(input_folder) as checks: + for check in checks: + if check.is_dir(): + # Remove imported checks + check_service = check.name.split("_")[0] + prowler_dir = prowler.__path__ + prowler_module = f"{prowler_dir[0]}/providers/{provider}/services/{check_service}/{check.name}" + if os.path.exists(prowler_module): + shutil.rmtree(prowler_module) + # If S3 URI, remove the downloaded folders + if s3_uri and os.path.exists(input_folder): + shutil.rmtree(input_folder) + + def list_services(provider: str) -> set(): available_services = set() checks_tuple = recover_checks_from_provider(provider) diff --git a/prowler/lib/cli/parser.py b/prowler/lib/cli/parser.py index 71778b52..89e7b4e4 100644 --- a/prowler/lib/cli/parser.py +++ b/prowler/lib/cli/parser.py @@ -226,6 +226,12 @@ Detailed documentation at https://docs.prowler.cloud default=[], # Pending validate choices ) + common_checks_parser.add_argument( + "-x", + "--checks-folder", + nargs="?", + help="Specify external directory with custom checks (each check must have a folder with the required files, see more in https://docs.prowler.cloud/en/latest/tutorials/misc/#custom-checks).", + ) def __init_list_checks_parser__(self): # List checks options diff --git a/prowler/lib/outputs/html.py b/prowler/lib/outputs/html.py index 3bc42f31..99b74d1d 100644 --- a/prowler/lib/outputs/html.py +++ b/prowler/lib/outputs/html.py @@ -211,31 +211,37 @@ def add_html_header(file_descriptor, audit_info): def fill_html(file_descriptor, finding, output_options): - row_class = "p-3 mb-2 bg-success-custom" - if finding.status == "INFO": - row_class = "table-info" - elif finding.status == "FAIL": - row_class = "table-danger" - elif finding.status == "WARNING": - row_class = "table-warning" - file_descriptor.write( - f""" - - {finding.status} - {finding.check_metadata.Severity} - {finding.check_metadata.ServiceName} - {finding.region} - {finding.check_metadata.CheckID.replace("_", "_")} - {finding.check_metadata.CheckTitle} - {finding.resource_id.replace("<", "<").replace(">", ">").replace("_", "_")} - {parse_html_string(unroll_tags(finding.resource_tags))} - {finding.status_extended.replace("<", "<").replace(">", ">").replace("_", "_")} -

{finding.check_metadata.Risk}

-

{finding.check_metadata.Remediation.Recommendation.Text}

-

{parse_html_string(unroll_dict(get_check_compliance(finding, finding.check_metadata.Provider, output_options)))}

- - """ - ) + try: + row_class = "p-3 mb-2 bg-success-custom" + if finding.status == "INFO": + row_class = "table-info" + elif finding.status == "FAIL": + row_class = "table-danger" + elif finding.status == "WARNING": + row_class = "table-warning" + file_descriptor.write( + f""" + + {finding.status} + {finding.check_metadata.Severity} + {finding.check_metadata.ServiceName} + {finding.region} + {finding.check_metadata.CheckID.replace("_", "_")} + {finding.check_metadata.CheckTitle} + {finding.resource_id.replace("<", "<").replace(">", ">").replace("_", "_")} + {parse_html_string(unroll_tags(finding.resource_tags))} + {finding.status_extended.replace("<", "<").replace(">", ">").replace("_", "_")} +

{finding.check_metadata.Risk}

+

{finding.check_metadata.Remediation.Recommendation.Text}

+

{parse_html_string(unroll_dict(get_check_compliance(finding, finding.check_metadata.Provider, output_options)))}

+ + """ + ) + except Exception as error: + logger.critical( + f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}] -- {error}" + ) + sys.exit(1) def fill_html_overview_statistics(stats, output_filename, output_directory): diff --git a/prowler/lib/outputs/models.py b/prowler/lib/outputs/models.py index 7d623a1b..635b9e37 100644 --- a/prowler/lib/outputs/models.py +++ b/prowler/lib/outputs/models.py @@ -12,20 +12,27 @@ from prowler.providers.aws.lib.audit_info.models import AWS_Organizations_Info def get_check_compliance(finding, provider, output_options): - check_compliance = {} - # We have to retrieve all the check's compliance requirements - for compliance in output_options.bulk_checks_metadata[ - finding.check_metadata.CheckID - ].Compliance: - compliance_fw = compliance.Framework - if compliance.Version: - compliance_fw = f"{compliance_fw}-{compliance.Version}" - if compliance.Provider == provider.upper(): - if compliance_fw not in check_compliance: - check_compliance[compliance_fw] = [] - for requirement in compliance.Requirements: - check_compliance[compliance_fw].append(requirement.Id) - return check_compliance + try: + check_compliance = {} + # We have to retrieve all the check's compliance requirements + if finding.check_metadata.CheckID in output_options.bulk_checks_metadata: + for compliance in output_options.bulk_checks_metadata[ + finding.check_metadata.CheckID + ].Compliance: + compliance_fw = compliance.Framework + if compliance.Version: + compliance_fw = f"{compliance_fw}-{compliance.Version}" + if compliance.Provider == provider.upper(): + if compliance_fw not in check_compliance: + check_compliance[compliance_fw] = [] + for requirement in compliance.Requirements: + check_compliance[compliance_fw].append(requirement.Id) + return check_compliance + except Exception as error: + logger.critical( + f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}] -- {error}" + ) + sys.exit(1) def generate_provider_output_csv( diff --git a/tests/lib/check/check_test.py b/tests/lib/check/check_test.py index 15fa4129..53041ecb 100644 --- a/tests/lib/check/check_test.py +++ b/tests/lib/check/check_test.py @@ -3,7 +3,9 @@ import pathlib from importlib.machinery import FileFinder from pkgutil import ModuleInfo +from boto3 import client, session from mock import patch +from moto import mock_s3 from prowler.lib.check.check import ( exclude_checks_to_run, @@ -11,8 +13,10 @@ from prowler.lib.check.check import ( list_modules, list_services, parse_checks_from_file, + parse_checks_from_folder, recover_checks_from_provider, recover_checks_from_service, + remove_custom_checks_module, update_audit_metadata, ) from prowler.lib.check.models import load_check_metadata @@ -20,6 +24,10 @@ from prowler.providers.aws.aws_provider import ( get_checks_from_input_arn, get_regions_from_audit_resources, ) +from prowler.providers.aws.lib.audit_info.models import AWS_Audit_Info + +AWS_ACCOUNT_NUMBER = 123456789012 +AWS_REGION = "us-east-1" expected_packages = [ ModuleInfo( @@ -127,6 +135,28 @@ def mock_recover_checks_from_aws_provider_lambda_service(*_): class Test_Check: + def set_mocked_audit_info(self): + audit_info = AWS_Audit_Info( + session_config=None, + original_session=None, + audit_session=session.Session( + profile_name=None, + botocore_session=None, + ), + audited_account=AWS_ACCOUNT_NUMBER, + audited_user_id=None, + audited_partition="aws", + audited_identity_arn=None, + profile=None, + profile_region=None, + credentials=None, + assumed_role_info=None, + audited_regions=None, + organizations_metadata=None, + audit_resources=None, + ) + return audit_info + def test_load_check_metadata(self): test_cases = [ { @@ -164,6 +194,50 @@ class Test_Check: provider = test["input"]["provider"] assert parse_checks_from_file(check_file, provider) == test["expected"] + @mock_s3 + def test_parse_checks_from_folder(self): + test_checks_folder = ( + f"{pathlib.Path().absolute()}/tests/lib/check/fixtures/checks_folder" + ) + # Create bucket and upload checks folder + s3_client = client("s3", region_name=AWS_REGION) + s3_client.create_bucket(Bucket="test") + # Iterate through the files in the folder and upload each one + for subdir, _, files in os.walk(test_checks_folder): + for file in files: + check = subdir.split("/")[-1] + full_path = os.path.join(subdir, file) + with open(full_path, "rb") as data: + s3_client.upload_fileobj( + data, "test", f"checks_folder/{check}/{file}" + ) + test_cases = [ + { + "input": { + "path": test_checks_folder, + "provider": "aws", + }, + "expected": 3, + }, + { + "input": { + "path": "s3://test/checks_folder/", + "provider": "aws", + }, + "expected": 3, + }, + ] + for test in test_cases: + check_folder = test["input"]["path"] + provider = test["input"]["provider"] + assert ( + parse_checks_from_folder( + self.set_mocked_audit_info(), check_folder, provider + ) + == test["expected"] + ) + remove_custom_checks_module(check_folder, provider) + def test_exclude_checks_to_run(self): test_cases = [ { diff --git a/tests/lib/check/fixtures/checks_folder/check11/__init__.py b/tests/lib/check/fixtures/checks_folder/check11/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/lib/check/fixtures/checks_folder/check12/__init__.py b/tests/lib/check/fixtures/checks_folder/check12/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/lib/check/fixtures/checks_folder/check7777/__init__.py b/tests/lib/check/fixtures/checks_folder/check7777/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/lib/cli/parser_test.py b/tests/lib/cli/parser_test.py index a171852b..17292539 100644 --- a/tests/lib/cli/parser_test.py +++ b/tests/lib/cli/parser_test.py @@ -31,6 +31,7 @@ class Test_Parser: assert not parsed.only_logs assert not parsed.checks assert not parsed.checks_file + assert not parsed.checks_folder assert not parsed.services assert not parsed.severity assert not parsed.compliance @@ -75,6 +76,7 @@ class Test_Parser: assert not parsed.only_logs assert not parsed.checks assert not parsed.checks_file + assert not parsed.checks_folder assert not parsed.services assert not parsed.severity assert not parsed.compliance @@ -111,6 +113,7 @@ class Test_Parser: assert not parsed.only_logs assert not parsed.checks assert not parsed.checks_file + assert not parsed.checks_folder assert not parsed.services assert not parsed.severity assert not parsed.compliance @@ -417,6 +420,20 @@ class Test_Parser: parsed = self.parser.parse(command) assert parsed.checks_file == filename + def test_checks_parser_checks_folder_short(self): + argument = "-x" + filename = "custom-checks-folder/" + command = [prowler_command, argument, filename] + parsed = self.parser.parse(command) + assert parsed.checks_folder == filename + + def test_checks_parser_checks_folder_long(self): + argument = "--checks-folder" + filename = "custom-checks-folder/" + command = [prowler_command, argument, filename] + parsed = self.parser.parse(command) + assert parsed.checks_folder == filename + def test_checks_parser_services_short(self): argument = "-s" service_1 = "iam" diff --git a/tests/lib/outputs/outputs_test.py b/tests/lib/outputs/outputs_test.py index edf93c1b..ad025dd5 100644 --- a/tests/lib/outputs/outputs_test.py +++ b/tests/lib/outputs/outputs_test.py @@ -807,6 +807,10 @@ class Test_Outputs: finding.status_extended = "This is a test" output_options = mock.MagicMock() + output_options.bulk_checks_metadata = {} + output_options.bulk_checks_metadata[ + "iam_disable_30_days_credentials" + ] = mock.MagicMock() output_options.bulk_checks_metadata[ "iam_disable_30_days_credentials" ].Compliance = bulk_check_metadata