From 66d57a3d36031bf57cc02d92074af79fbd90bf90 Mon Sep 17 00:00:00 2001
From: Pepe Fagoaga <pepe@verica.io>
Date: Thu, 23 Jun 2022 16:56:06 +0200
Subject: [PATCH] feat(severity): Run checks by severity (#1223)

---
 lib/check/check.py         | 14 ++++++++++++++
 lib/check/checks_loader.py | 12 ++++++++----
 prowler                    | 33 +++++++++++++++++++--------------
 3 files changed, 41 insertions(+), 18 deletions(-)

diff --git a/lib/check/check.py b/lib/check/check.py
index 39f64baa..7bda1ffc 100644
--- a/lib/check/check.py
+++ b/lib/check/check.py
@@ -98,6 +98,20 @@ def print_services(service_list: set):
         print(f"- {service}")
 
 
+
+def print_checks(provider: str, check_list: set, bulk_checks_metadata: dict):
+    for check in check_list:
+        try:
+            print(
+                f"[{bulk_checks_metadata[check].CheckID}] {bulk_checks_metadata[check].CheckTitle} - {Fore.MAGENTA}{bulk_checks_metadata[check].ServiceName} {Fore.YELLOW}[{bulk_checks_metadata[check].Severity}]{Style.RESET_ALL}"
+            )
+        except KeyError as error:
+            logger.error(
+                f"Check {error} was not found for the {provider.upper()} provider"
+            )
+
+
+
 # List available groups
 def list_groups(provider: str):
     groups = parse_groups_from_file(groups_file)
diff --git a/lib/check/checks_loader.py b/lib/check/checks_loader.py
index 80c2e342..cf6edb29 100644
--- a/lib/check/checks_loader.py
+++ b/lib/check/checks_loader.py
@@ -15,6 +15,7 @@ def load_checks_to_execute(
     check_list: list,
     service_list: list,
     group_list: list,
+    severities: list,
     provider: str,
 ) -> set:
 
@@ -25,10 +26,13 @@ def load_checks_to_execute(
         for check_name in check_list:
             checks_to_execute.add(check_name)
 
-    # elif severity_list:
-    # using bulk_checks_metadata
-    # elif compliance_list:
-    # using bulk_checks_metadata
+    # Handle if there are some severities passed using --severity
+    elif severities:
+        for check in bulk_checks_metadata:
+            # Check check's severity
+            if bulk_checks_metadata[check].Severity in severities:
+                checks_to_execute.add(check)
+
     # Handle if there are checks passed using -C/--checks-file
     elif checks_file:
         try:
diff --git a/prowler b/prowler
index 0a826bc1..fe6bfa01 100755
--- a/prowler
+++ b/prowler
@@ -4,8 +4,6 @@
 import argparse
 import sys
 
-from colorama import Fore, Style
-
 from lib.banner import print_banner, print_version
 from lib.check.check import (
     bulk_load_checks_metadata,
@@ -15,6 +13,7 @@ from lib.check.check import (
     import_check,
     list_groups,
     list_services,
+    print_checks,
     print_services,
     run_check,
     set_output_options,
@@ -35,12 +34,19 @@ if __name__ == "__main__":
     group.add_argument("-C", "--checks-file", nargs="?", help="List of checks")
     group.add_argument("-s", "--services", nargs="+", help="List of services")
     group.add_argument("-g", "--groups", nargs="+", help="List of groups")
+    group.add_argument(
+        "--severity",
+        nargs="+",
+        help="List of severities [informational, low, medium, high, critical]",
+        choices=["informational","low","medium","high","critical"]
+    )
+    # Exclude checks options
     parser.add_argument("-e", "--excluded-checks", nargs="+", help="Checks to exclude")
     parser.add_argument("-E", "--excluded-groups", nargs="+", help="Groups to exclude")
     parser.add_argument(
         "-S", "--excluded-services", nargs="+", help="Services to exclude"
     )
-
+    # List checks options
     list_group = parser.add_mutually_exclusive_group()
     list_group.add_argument(
         "-L", "--list-groups", action="store_true", help="List groups"
@@ -53,7 +59,7 @@ if __name__ == "__main__":
     )
 
     parser.add_argument(
-        "-b", "--no-banner", action="store_false", help="Hide Prowler Banner"
+        "-b", "--no-banner", action="store_false", help="Hide Prowler banner"
     )
     parser.add_argument(
         "-v", "--version", action="store_true", help="Show Prowler version"
@@ -121,6 +127,7 @@ if __name__ == "__main__":
     services = args.services
     groups = args.groups
     checks_file = args.checks_file
+    severities = args.severity
 
     # Set Logger configuration
     set_logging_config(args.log_file, args.log_level)
@@ -155,7 +162,13 @@ if __name__ == "__main__":
 
     # Load checks to execute
     checks_to_execute = load_checks_to_execute(
-        bulk_checks_metadata, checks_file, checks, services, groups, provider
+        bulk_checks_metadata,
+        checks_file,
+        checks,
+        services,
+        groups,
+        severities,
+        provider,
     )
     # Exclude checks if -e/--excluded-checks
     if excluded_checks:
@@ -175,15 +188,7 @@ if __name__ == "__main__":
 
     # If -l/--list-checks passed as argument, print checks to execute and quit
     if args.list_checks:
-        for check in checks_to_execute:
-            try:
-                print(
-                    f"[{bulk_checks_metadata[check].CheckID}] {bulk_checks_metadata[check].CheckTitle} - {Fore.MAGENTA}{bulk_checks_metadata[check].ServiceName} {Fore.YELLOW}[{bulk_checks_metadata[check].Severity}]{Style.RESET_ALL}"
-                )
-            except KeyError as error:
-                logger.error(
-                    f"Check {error} was not found for the {provider.upper()} provider"
-                )
+        print_checks(provider, checks_to_execute, bulk_checks_metadata)
         sys.exit()
 
     # Setting output options