refactor(execute_check): refactor execute method (#4975)

prowler/__main__.py

@@ -455,8 +455,8 @@ def prowler():
             elif compliance_name.startswith("kisa"):
                 # Generate KISA-ISMS-P Finding Object
                 filename = (
-                    f"{global_provider.output_options.output_directory}/compliance/"
-                    f"{global_provider.output_options.output_filename}_{compliance_name}.csv"
+                    f"{output_options.output_directory}/compliance/"
+                    f"{output_options.output_filename}_{compliance_name}.csv"
                 )
                 kisa_ismsp = AWSKISAISMSP(
                     findings=finding_outputs,

prowler/lib/check/check.py

@@ -325,35 +325,6 @@ def import_check(check_path: str) -> ModuleType:
     return lib
 
 
-def run_check(check: Check, verbose: bool = False, only_logs: bool = False) -> list:
-    """
-    Run the check and return the findings
-    Args:
-        check (Check): check class
-        output_options (Any): output options
-    Returns:
-        list: list of findings
-    """
-    findings = []
-    if verbose:
-        print(
-            f"\nCheck ID: {check.CheckID} - {Fore.MAGENTA}{check.ServiceName}{Fore.YELLOW} [{check.Severity}]{Style.RESET_ALL}"
-        )
-    logger.debug(f"Executing check: {check.CheckID}")
-    try:
-        findings = check.execute()
-    except Exception as error:
-        if not only_logs:
-            print(
-                f"Something went wrong in {check.CheckID}, please use --log-level ERROR"
-            )
-        logger.error(
-            f"{check.CheckID} -- {error.__class__.__name__}[{traceback.extract_tb(error.__traceback__)[-1].lineno}]: {error}"
-        )
-    finally:
-        return findings
-
-
 def run_fixer(check_findings: list) -> int:
     """
     Run the fixer for the check if it exists and there are any FAIL findings
@@ -471,19 +442,42 @@ def execute_checks(
                 f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
             )
 
+    # Set verbose flag
+    verbose = False
+    if hasattr(output_options, "verbose"):
+        verbose = output_options.verbose
+    elif hasattr(output_options, "fixer"):
+        verbose = output_options.fixer
+
     # Execution with the --only-logs flag
     if output_options.only_logs:
         for check_name in checks_to_execute:
             # Recover service from check name
             service = check_name.split("_")[0]
             try:
+                try:
+                    # Import check module
+                    check_module_path = f"prowler.providers.{global_provider.type}.services.{service}.{check_name}.{check_name}"
+                    lib = import_check(check_module_path)
+                    # Recover functions from check
+                    check_to_execute = getattr(lib, check_name)
+                    check = check_to_execute()
+                except ModuleNotFoundError:
+                    logger.error(
+                        f"Check '{check_name}' was not found for the {global_provider.type.upper()} provider"
+                    )
+                    continue
+                if verbose:
+                    print(
+                        f"\nCheck ID: {check.CheckID} - {Fore.MAGENTA}{check.ServiceName}{Fore.YELLOW} [{check.Severity}]{Style.RESET_ALL}"
+                    )
                 check_findings = execute(
-                    service,
-                    check_name,
+                    check,
                     global_provider,
                     custom_checks_metadata,
                     output_options,
                 )
+                report(check_findings, global_provider, output_options)
                 all_findings.extend(check_findings)
 
                 # Update Audit Status
@@ -541,13 +535,31 @@ def execute_checks(
                     f"-> Scanning {orange_color}{service}{Style.RESET_ALL} service"
                 )
                 try:
+                    try:
+                        # Import check module
+                        check_module_path = f"prowler.providers.{global_provider.type}.services.{service}.{check_name}.{check_name}"
+                        lib = import_check(check_module_path)
+                        # Recover functions from check
+                        check_to_execute = getattr(lib, check_name)
+                        check = check_to_execute()
+                    except ModuleNotFoundError:
+                        logger.error(
+                            f"Check '{check_name}' was not found for the {global_provider.type.upper()} provider"
+                        )
+                        continue
+                    if verbose:
+                        print(
+                            f"\nCheck ID: {check.CheckID} - {Fore.MAGENTA}{check.ServiceName}{Fore.YELLOW} [{check.Severity}]{Style.RESET_ALL}"
+                        )
                     check_findings = execute(
-                        service,
-                        check_name,
+                        check,
                         global_provider,
                         custom_checks_metadata,
                         output_options,
                     )
+
+                    report(check_findings, global_provider, output_options)
+
                     all_findings.extend(check_findings)
                     services_executed.add(service)
                     checks_executed.add(check_name)
@@ -570,12 +582,25 @@ def execute_checks(
                     )
                 bar()
             bar.title = f"-> {Fore.GREEN}Scan completed!{Style.RESET_ALL}"
+
+    # Custom report interface
+    if os.environ.get("PROWLER_REPORT_LIB_PATH"):
+        try:
+            logger.info("Using custom report interface ...")
+            lib = os.environ["PROWLER_REPORT_LIB_PATH"]
+            outputs_module = importlib.import_module(lib)
+            custom_report_interface = getattr(outputs_module, "report")
+
+            # TODO: review this call and see if we can remove the global_provider.output_options since it is contained in the global_provider
+            custom_report_interface(check_findings, output_options, global_provider)
+        except Exception:
+            sys.exit(1)
+
     return all_findings
 
 
 def execute(
-    service: str,
-    check_name: str,
+    check: Check,
     global_provider: Any,
     custom_checks_metadata: Any,
     output_options: Any = None,
@@ -594,33 +619,31 @@ def execute(
         list: list of findings
     """
     try:
-        # Import check module
-        check_module_path = f"prowler.providers.{global_provider.type}.services.{service}.{check_name}.{check_name}"
-        lib = import_check(check_module_path)
-        # Recover functions from check
-        check_to_execute = getattr(lib, check_name)
-        check_class = check_to_execute()
-
         # Update check metadata to reflect that in the outputs
         if custom_checks_metadata and custom_checks_metadata["Checks"].get(
-            check_class.CheckID
+            check.CheckID
         ):
-            check_class = update_check_metadata(
-                check_class, custom_checks_metadata["Checks"][check_class.CheckID]
+            check = update_check_metadata(
+                check, custom_checks_metadata["Checks"][check.CheckID]
             )
 
-        # Run check
-        verbose = False
-        if hasattr(output_options, "verbose"):
-            verbose = output_options.verbose
-        elif hasattr(output_options, "fixer"):
-            verbose = output_options.fixer
-
         only_logs = False
         if hasattr(output_options, "only_logs"):
             only_logs = output_options.only_logs
 
-        check_findings = run_check(check_class, verbose, only_logs)
+        # Execute the check
+        check_findings = []
+        logger.debug(f"Executing check: {check.CheckID}")
+        try:
+            check_findings = check.execute()
+        except Exception as error:
+            if not only_logs:
+                print(
+                    f"Something went wrong in {check.CheckID}, please use --log-level ERROR"
+                )
+            logger.error(
+                f"{check.CheckID} -- {error.__class__.__name__}[{traceback.extract_tb(error.__traceback__)[-1].lineno}]: {error}"
+            )
 
         # Exclude findings per status
         if hasattr(output_options, "status") and output_options.status:
@@ -630,9 +653,8 @@ def execute(
                 if finding.status in output_options.status
             ]
 
-        # Mutelist findings
+        # Before returning the findings, we need to apply the mute list logic
         if hasattr(global_provider, "mutelist") and global_provider.mutelist.mutelist:
-            # TODO: make this prettier
             is_finding_muted_args = {}
             if global_provider.type == "aws":
                 is_finding_muted_args["aws_account_id"] = (
@@ -647,25 +669,9 @@ def execute(
                     **is_finding_muted_args
                 )
 
-        # Refactor(Outputs)
-        # Report the check's findings
-        report(check_findings, global_provider, output_options)
-
-        # Refactor(Outputs)
-        if os.environ.get("PROWLER_REPORT_LIB_PATH"):
-            try:
-                logger.info("Using custom report interface ...")
-                lib = os.environ["PROWLER_REPORT_LIB_PATH"]
-                outputs_module = importlib.import_module(lib)
-                custom_report_interface = getattr(outputs_module, "report")
-
-                # TODO: review this call and see if we can remove the output_options since it is contained in the global_provider
-                custom_report_interface(check_findings, output_options, global_provider)
-            except Exception:
-                sys.exit(1)
     except ModuleNotFoundError:
         logger.error(
-            f"Check '{check_name}' was not found for the {global_provider.type.upper()} provider"
+            f"Check '{check.CheckID}' was not found for the {global_provider.type.upper()} provider"
         )
         check_findings = []
     except Exception as error:

prowler/lib/scan/scan.py

@@ -1,6 +1,6 @@
 from typing import Generator
 
-from prowler.lib.check.check import execute, update_audit_metadata
+from prowler.lib.check.check import execute, import_check, update_audit_metadata
 from prowler.lib.logger import logger
 from prowler.lib.outputs.finding import Finding
 from prowler.providers.common.models import Audit_Metadata
@@ -99,11 +99,21 @@ class Scan:
                 try:
                     # Recover service from check name
                     service = get_service_name_from_check_name(check_name)
-
+                    try:
+                        # Import check module
+                        check_module_path = f"prowler.providers.{self._provider.type}.services.{service}.{check_name}.{check_name}"
+                        lib = import_check(check_module_path)
+                        # Recover functions from check
+                        check_to_execute = getattr(lib, check_name)
+                        check = check_to_execute()
+                    except ModuleNotFoundError:
+                        logger.error(
+                            f"Check '{check_name}' was not found for the {self._provider.type.upper()} provider"
+                        )
+                        continue
                     # Execute the check
                     check_findings = execute(
-                        service,
-                        check_name,
+                        check,
                         self._provider,
                         custom_checks_metadata,
                         output_options=None,

tests/lib/check/check_test.py

@@ -1,14 +1,12 @@
 import json
 import os
 import pathlib
-import traceback
 from importlib.machinery import FileFinder
-from logging import DEBUG, ERROR
+from logging import ERROR
 from pkgutil import ModuleInfo
 from unittest import mock
 
 from boto3 import client
-from colorama import Fore, Style
 from mock import Mock, patch
 from moto import mock_aws
 
@@ -16,13 +14,13 @@ from prowler.lib.check.check import (
     exclude_checks_to_run,
     exclude_services_to_run,
     execute,
+    execute_checks,
     list_categories,
     list_checks_json,
     list_services,
     parse_checks_from_file,
     parse_checks_from_folder,
     remove_custom_checks_module,
-    run_check,
     update_audit_metadata,
 )
 from prowler.lib.check.models import load_check_metadata
@@ -815,13 +813,19 @@ class TestCheck:
                 region=AWS_REGION_US_EAST_1,
             )
         ]
-        with mock.patch(
+        finding = Mock()
+        finding.status = "PASS"
+        findings = [finding]
+        check = Mock()
+        check.CheckID = "test-check"
+        check.execute = Mock(return_value=findings)
+
+        with patch("prowler.lib.check.check.execute", return_value=findings), patch(
             "prowler.providers.aws.services.accessanalyzer.accessanalyzer_service.AccessAnalyzer",
             accessanalyzer_client,
         ):
             findings = execute(
-                service="accessanalyzer",
-                check_name="accessanalyzer_enabled",
+                check=check,
                 global_provider=set_mocked_aws_provider(
                     expected_checks=["accessanalyzer_enabled"]
                 ),
@@ -844,6 +848,9 @@ class TestCheck:
             )
         ]
         status = ["PASS"]
+        check = mock.MagicMock()
+        check.CheckID = "accessanalyzer_enabled"
+        check.Status = "PASS"
         output_options = mock.MagicMock()
         output_options.status = status
         with mock.patch(
@@ -851,8 +858,7 @@ class TestCheck:
             accessanalyzer_client,
         ):
             findings = execute(
-                service="accessanalyzer",
-                check_name="accessanalyzer_enabled",
+                check=check,
                 global_provider=set_mocked_aws_provider(
                     expected_checks=["accessanalyzer_enabled"]
                 ),
@@ -861,92 +867,6 @@ class TestCheck:
             )
             assert len(findings) == 0
 
-    def test_run_check(self, caplog):
-        caplog.set_level(DEBUG)
-
-        findings = []
-        check = Mock()
-        check.CheckID = "test-check"
-        check.execute = Mock(return_value=findings)
-
-        with patch("prowler.lib.check.check.execute", return_value=findings):
-            assert run_check(check) == findings
-            assert caplog.record_tuples == [
-                (
-                    "root",
-                    DEBUG,
-                    f"Executing check: {check.CheckID}",
-                )
-            ]
-
-    def test_run_check_verbose(self, capsys):
-
-        findings = []
-        check = Mock()
-        check.CheckID = "test-check"
-        check.ServiceName = "test-service"
-        check.Severity = "test-severity"
-        check.execute = Mock(return_value=findings)
-
-        with patch("prowler.lib.check.check.execute", return_value=findings):
-            assert run_check(check, verbose=True) == findings
-            assert (
-                capsys.readouterr().out
-                == f"\nCheck ID: {check.CheckID} - {Fore.MAGENTA}{check.ServiceName}{Fore.YELLOW} [{check.Severity}]{Style.RESET_ALL}\n"
-            )
-
-    def test_run_check_exception_only_logs(self, caplog):
-        caplog.set_level(ERROR)
-
-        findings = []
-        check = Mock()
-        check.CheckID = "test-check"
-        check.ServiceName = "test-service"
-        check.Severity = "test-severity"
-        error = Exception()
-        check.execute = Mock(side_effect=error)
-
-        with patch("prowler.lib.check.check.execute", return_value=findings):
-            assert run_check(check, only_logs=True) == findings
-            assert caplog.record_tuples == [
-                (
-                    "root",
-                    ERROR,
-                    f"{check.CheckID} -- {error.__class__.__name__}[{traceback.extract_tb(error.__traceback__)[-1].lineno}]: {error}",
-                )
-            ]
-
-    def test_run_check_exception(self, caplog, capsys):
-        caplog.set_level(ERROR)
-
-        findings = []
-        check = Mock()
-        check.CheckID = "test-check"
-        check.ServiceName = "test-service"
-        check.Severity = "test-severity"
-        error = Exception()
-        check.execute = Mock(side_effect=error)
-
-        with patch("prowler.lib.check.check.execute", return_value=findings):
-            assert (
-                run_check(
-                    check,
-                    verbose=False,
-                )
-                == findings
-            )
-            assert caplog.record_tuples == [
-                (
-                    "root",
-                    ERROR,
-                    f"{check.CheckID} -- {error.__class__.__name__}[{traceback.extract_tb(error.__traceback__)[-1].lineno}]: {error}",
-                )
-            ]
-            assert (
-                capsys.readouterr().out
-                == f"Something went wrong in {check.CheckID}, please use --log-level ERROR\n"
-            )
-
     def test_aws_checks_metadata_is_valid(self):
         # Check if the checkID in the metadata.json of the checks is correct
         # Define the base directory for the checks
@@ -1020,3 +940,33 @@ class TestCheck:
                                 assert (
                                     check_id == check_dir
                                 ), f"CheckID in metadata does not match the check name in {check_directory}. Found CheckID: {check_id}"
+
+    def test_execute_check_exception_only_logs(self, caplog):
+        caplog.set_level(ERROR)
+
+        findings = []
+        check = Mock()
+        checks = ["test-check"]
+
+        provider = mock.MagicMock()
+        provider.type = "aws"
+
+        output_options = mock.MagicMock()
+        output_options.only_logs = True
+        error = Exception()
+        check.execute = Mock(side_effect=error)
+
+        with patch("prowler.lib.check.check.execute", return_value=findings):
+            assert (
+                execute_checks(
+                    checks,
+                    provider,
+                    custom_checks_metadata=None,
+                    config_file=None,
+                    output_options=output_options,
+                )
+                == findings
+            )
+            assert caplog.record_tuples == [
+                ("root", 40, f"Check '{checks[0]}' was not found for the AWS provider")
+            ]

tests/lib/outputs/output_options/output_options_test.py

@@ -1,5 +1,6 @@
 from argparse import Namespace
 from datetime import datetime
+from os import rmdir
 from unittest import mock
 
 from freezegun import freeze_time
@@ -45,6 +46,9 @@ class Test_Output_Options:
         )
         assert output_options.bulk_checks_metadata == {}
 
+        rmdir(f"{arguments.output_directory}/compliance")
+        rmdir(arguments.output_directory)
+
     @freeze_time(datetime.today())
     def test_set_output_options_aws(self):
         arguments = Namespace()
@@ -74,6 +78,9 @@ class Test_Output_Options:
         assert output_options.verbose
         assert output_options.output_filename == arguments.output_filename
 
+        rmdir(f"{arguments.output_directory}/compliance")
+        rmdir(arguments.output_directory)
+
     @freeze_time(datetime.today())
     def test_azure_provider_output_options_with_domain(self):
         arguments = Namespace()
@@ -109,6 +116,9 @@ class Test_Output_Options:
             == f"prowler-output-{identity.tenant_domain}-{output_file_timestamp}"
         )
 
+        rmdir(f"{arguments.output_directory}/compliance")
+        rmdir(arguments.output_directory)
+
     @freeze_time(datetime.today())
     def test_gcp_output_options(self):
         arguments = Namespace()
@@ -140,6 +150,9 @@ class Test_Output_Options:
         assert output_optionss.verbose
         assert f"prowler-output-{identity.profile}" in output_optionss.output_filename
 
+        rmdir(f"{arguments.output_directory}/compliance")
+        rmdir(arguments.output_directory)
+
     def test_set_output_options_kubernetes(self):
         arguments = Namespace()
         arguments.status = []
@@ -167,3 +180,6 @@ class Test_Output_Options:
         assert output_options.bulk_checks_metadata == {}
         assert output_options.verbose
         assert output_options.output_filename == arguments.output_filename
+
+        rmdir(f"{arguments.output_directory}/compliance")
+        rmdir(arguments.output_directory)

tests/lib/scan/scan_test.py

@@ -1,6 +1,7 @@
 from unittest import mock
 
 import pytest
+from mock import MagicMock, patch
 
 from prowler.lib.scan.scan import Scan, get_service_checks_to_execute
 from tests.lib.outputs.fixtures.fixtures import generate_finding_output
@@ -203,42 +204,40 @@ class TestScan:
         assert scan.get_completed_services() == set()
         assert scan.get_completed_checks() == set()
 
+    @patch("importlib.import_module")
     def test_scan(
-        mock_global_provider, mock_execute, mock_logger, mock_generate_output
+        mock_import_module,
+        mock_global_provider,
+        mock_execute,
+        mock_logger,
+        mock_generate_output,
     ):
-        checks_to_execute = {"accessanalyzer_enabled", "ec2_instance_public"}
+        mock_check_class = MagicMock()
+        mock_check_instance = mock_check_class.return_value
+        mock_check_instance.Provider = "aws"
+        mock_check_instance.CheckID = "accessanalyzer_enabled"
+        mock_check_instance.CheckTitle = "Check if IAM Access Analyzer is enabled"
+
+        mock_import_module.return_value = MagicMock(
+            accessanalyzer_enabled=mock_check_class
+        )
+
+        checks_to_execute = {"accessanalyzer_enabled"}
         custom_checks_metadata = {}
+        mock_global_provider.type = "aws"
 
-        # Create a Scan object
         scan = Scan(mock_global_provider, checks_to_execute)
-
-        # Execute the scan
         results = list(scan.scan(custom_checks_metadata))
 
-        # Verify that generate_output was called with the correct findings
-        assert mock_generate_output.call_count == 2 * len(mock_execute.side_effect())
-
-        # Verify that execute was called twice
-        assert mock_execute.call_count == 2
-
-        assert len(results) == 2
+        assert mock_generate_output.call_count == 1 * len(mock_execute.side_effect())
+        assert mock_execute.call_count == 1
+        assert len(results) == 1
         assert results[0][1] == mock_execute.side_effect()
-        assert results[1][1] == mock_execute.side_effect()
-
-        # Check the audit progress for the last result
-        assert results[1][0] == 100.0
-
-        # Verify that the progress is 100.0
-        assert scan.progress == 100.0  # 100% progress is 100
-        assert scan._number_of_checks_completed == 2
-        assert scan.service_checks_to_execute == {}
+        assert results[0][0] == 100.0
+        assert scan.progress == 100.0
+        assert scan._number_of_checks_completed == 1
         assert scan.service_checks_completed == {
-            "ec2": {"ec2_instance_public"},
             "accessanalyzer": {"accessanalyzer_enabled"},
         }
-
-        # Verify that the findings are correct
-        assert scan.findings == mock_execute.side_effect() + mock_execute.side_effect()
-
-        # Verify that no error was logged
+        assert scan.findings == mock_execute.side_effect()
         mock_logger.error.assert_not_called()