From 2a139e3dc7068a2b34e3e30e596c6fc62424934a Mon Sep 17 00:00:00 2001
From: Seiji Ujihira <sejimhp@gmail.com>
Date: Fri, 7 Jun 2024 23:01:28 +0900
Subject: [PATCH] fix(custom): execute custom checks (#4202)

---
 prowler/__main__.py           | 3 ++-
 prowler/lib/check/check.py    | 8 ++++----
 tests/lib/check/check_test.py | 4 ++--
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/prowler/__main__.py b/prowler/__main__.py
index d8ab4061aa..9e99663410 100644
--- a/prowler/__main__.py
+++ b/prowler/__main__.py
@@ -180,7 +180,8 @@ def prowler():
 
     # Import custom checks from folder
     if checks_folder:
-        parse_checks_from_folder(global_provider, checks_folder)
+        custom_checks = parse_checks_from_folder(global_provider, checks_folder)
+        checks_to_execute.update(custom_checks)
 
     # Exclude checks if -e/--excluded-checks
     if excluded_checks:
diff --git a/prowler/lib/check/check.py b/prowler/lib/check/check.py
index b6592e5ee9..c7545eb48e 100644
--- a/prowler/lib/check/check.py
+++ b/prowler/lib/check/check.py
@@ -126,9 +126,9 @@ def parse_checks_from_file(input_file: str, provider: str) -> set:
 
 
 # Load checks from custom folder
-def parse_checks_from_folder(provider, input_folder: str) -> int:
+def parse_checks_from_folder(provider, input_folder: str) -> set:
     try:
-        imported_checks = 0
+        custom_checks = set()
         # Check if input folder is a S3 URI
         if provider.type == "aws" and re.search(
             "^s3://([^/]+)/(.*?([^/]+))/$", input_folder
@@ -156,8 +156,8 @@ def parse_checks_from_folder(provider, input_folder: str) -> int:
                     if os.path.exists(prowler_module):
                         shutil.rmtree(prowler_module)
                     shutil.copytree(check_module, prowler_module)
-                    imported_checks += 1
-        return imported_checks
+                    custom_checks.add(check.name)
+        return custom_checks
     except Exception as error:
         logger.critical(
             f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}] -- {error}"
diff --git a/tests/lib/check/check_test.py b/tests/lib/check/check_test.py
index 6cb7e05cdf..32576f0a5b 100644
--- a/tests/lib/check/check_test.py
+++ b/tests/lib/check/check_test.py
@@ -453,14 +453,14 @@ class TestCheck:
                     "path": test_checks_folder,
                     "provider": "aws",
                 },
-                "expected": 3,
+                "expected": {"check11", "check12", "check7777"},
             },
             {
                 "input": {
                     "path": "s3://test/checks_folder/",
                     "provider": "aws",
                 },
-                "expected": 3,
+                "expected": {"check11", "check12", "check7777"},
             },
         ]