feat(custom_filename): custom output filename (#1345)

* feat(s3_output): send outputs to S3 bucket

* feat(custom_filename): custom output filename

Co-authored-by: sergargar <sergio@verica.io>

lib/outputs/outputs.py

@@ -1,5 +1,6 @@
 import json
 import os
+import sys
 from csv import DictWriter
 
 from colorama import Fore, Style
@@ -12,6 +13,7 @@ from config.config import (
     timestamp_iso,
     timestamp_utc,
 )
+from lib.logger import logger
 from lib.outputs.models import (
     Check_Output_CSV,
     Check_Output_JSON,
@@ -37,9 +39,9 @@ def report(check_findings, output_options, audit_info):
 
         file_descriptors = fill_file_descriptors(
             output_options.output_modes,
-            audit_info.audited_account,
             output_options.output_directory,
             csv_fields,
+            output_options.output_filename,
         )
 
     if check_findings:
@@ -101,13 +103,11 @@ def report(check_findings, output_options, audit_info):
             file_descriptors.get(file_descriptor).close()
 
 
-def fill_file_descriptors(output_modes, audited_account, output_directory, csv_fields):
+def fill_file_descriptors(output_modes, output_directory, csv_fields, output_filename):
     file_descriptors = {}
     for output_mode in output_modes:
         if output_mode == "csv":
-            filename = (
-                f"{output_directory}/prowler-output-{audited_account}-{csv_file_suffix}"
-            )
+            filename = f"{output_directory}/{output_filename}{csv_file_suffix}"
             if file_exists(filename):
                 file_descriptor = open_file(
                     filename,
@@ -127,7 +127,7 @@ def fill_file_descriptors(output_modes, audited_account, output_directory, csv_f
             file_descriptors.update({output_mode: file_descriptor})
 
         if output_mode == "json":
-            filename = f"{output_directory}/prowler-output-{audited_account}-{json_file_suffix}"
+            filename = f"{output_directory}/{output_filename}{json_file_suffix}"
             if file_exists(filename):
                 file_descriptor = open_file(
                     filename,
@@ -143,7 +143,7 @@ def fill_file_descriptors(output_modes, audited_account, output_directory, csv_f
             file_descriptors.update({output_mode: file_descriptor})
 
         if output_mode == "json-asff":
-            filename = f"{output_directory}/prowler-output-{audited_account}-{json_asff_file_suffix}"
+            filename = f"{output_directory}/{output_filename}{json_asff_file_suffix}"
             if file_exists(filename):
                 file_descriptor = open_file(
                     filename,
@@ -238,17 +238,47 @@ def fill_json_asff(finding_output, audit_info, finding):
     return finding_output
 
 
-def close_json(output_directory, audited_account, mode):
-    suffix = json_file_suffix
-    if mode == "json-asff":
-        suffix = json_asff_file_suffix
-    filename = f"{output_directory}/prowler-output-{audited_account}-{suffix}"
-    file_descriptor = open_file(
-        filename,
-        "a",
-    )
-    # Replace last comma for square bracket
-    file_descriptor.seek(file_descriptor.tell() - 1, os.SEEK_SET)
-    file_descriptor.truncate()
-    file_descriptor.write("]")
-    file_descriptor.close()
+def close_json(output_filename, output_directory, mode):
+    try:
+        suffix = json_file_suffix
+        if mode == "json-asff":
+            suffix = json_asff_file_suffix
+        filename = f"{output_directory}/{output_filename}{suffix}"
+        file_descriptor = open_file(
+            filename,
+            "a",
+        )
+        # Replace last comma for square bracket
+        file_descriptor.seek(file_descriptor.tell() - 1, os.SEEK_SET)
+        file_descriptor.truncate()
+        file_descriptor.write("]")
+        file_descriptor.close()
+    except Exception as error:
+        logger.critical(f"{error.__class__.__name__} -- {error}")
+        sys.exit()
+
+
+def send_to_s3_bucket(
+    output_filename, output_directory, output_mode, output_bucket, audit_session
+):
+    try:
+        # Get only last part of the path
+        output_directory = output_directory.split("/")[-1]
+        if output_mode == "csv":
+            filename = f"{output_filename}{csv_file_suffix}"
+        elif output_mode == "json":
+            filename = f"{output_filename}{json_file_suffix}"
+        elif output_mode == "json-asff":
+            filename = f"{output_filename}{json_asff_file_suffix}"
+        logger.info(f"Sending outputs to S3 bucket {output_bucket}")
+        # Check if security hub is enabled in current region
+        s3_client = audit_session.client("s3")
+        s3_client.upload_file(
+            output_directory + "/" + filename,
+            output_bucket,
+            output_directory + "/" + output_mode + "/" + filename,
+        )
+
+    except Exception as error:
+        logger.critical(f"{error.__class__.__name__} -- {error}")
+        sys.exit()

lib/outputs/outputs_test.py

@@ -1,3 +1,4 @@
+import os
 from os import path, remove
 
 from colorama import Fore
@@ -6,6 +7,7 @@ from config.config import (
     csv_file_suffix,
     json_asff_file_suffix,
     json_file_suffix,
+    output_file_timestamp,
     prowler_version,
     timestamp_iso,
     timestamp_utc,
@@ -33,7 +35,7 @@ from providers.aws.models import AWS_Audit_Info
 class Test_Outputs:
     def test_fill_file_descriptors(self):
         audited_account = "123456789012"
-        output_directory = "."
+        output_directory = f"{os.path.dirname(os.path.realpath(__file__))}"
         csv_fields = generate_csv_fields()
         test_output_modes = [
             ["csv"],
@@ -42,47 +44,47 @@ class Test_Outputs:
             ["csv", "json"],
             ["csv", "json", "json-asff"],
         ]
-
+        output_filename = f"prowler-output-{audited_account}-{output_file_timestamp}"
         expected = [
             {
                 "csv": open_file(
-                    f"{output_directory}/prowler-output-{audited_account}-{csv_file_suffix}",
+                    f"{output_directory}/{output_filename}{csv_file_suffix}",
                     "a",
                 )
             },
             {
                 "json": open_file(
-                    f"{output_directory}/prowler-output-{audited_account}-{json_file_suffix}",
+                    f"{output_directory}/{output_filename}{json_file_suffix}",
                     "a",
                 )
             },
             {
                 "json-asff": open_file(
-                    f"{output_directory}/prowler-output-{audited_account}-{json_asff_file_suffix}",
+                    f"{output_directory}/{output_filename}{json_asff_file_suffix}",
                     "a",
                 )
             },
             {
                 "csv": open_file(
-                    f"{output_directory}/prowler-output-{audited_account}-{csv_file_suffix}",
+                    f"{output_directory}/{output_filename}{csv_file_suffix}",
                     "a",
                 ),
                 "json": open_file(
-                    f"{output_directory}/prowler-output-{audited_account}-{json_file_suffix}",
+                    f"{output_directory}/{output_filename}{json_file_suffix}",
                     "a",
                 ),
             },
             {
                 "csv": open_file(
-                    f"{output_directory}/prowler-output-{audited_account}-{csv_file_suffix}",
+                    f"{output_directory}/{output_filename}{csv_file_suffix}",
                     "a",
                 ),
                 "json": open_file(
-                    f"{output_directory}/prowler-output-{audited_account}-{json_file_suffix}",
+                    f"{output_directory}/{output_filename}{json_file_suffix}",
                     "a",
                 ),
                 "json-asff": open_file(
-                    f"{output_directory}/prowler-output-{audited_account}-{json_asff_file_suffix}",
+                    f"{output_directory}/{output_filename}{json_asff_file_suffix}",
                     "a",
                 ),
             },
@@ -90,7 +92,10 @@ class Test_Outputs:
 
         for index, output_mode_list in enumerate(test_output_modes):
             test_output_file_descriptors = fill_file_descriptors(
-                output_mode_list, audited_account, output_directory, csv_fields
+                output_mode_list,
+                output_directory,
+                csv_fields,
+                output_filename,
             )
             for output_mode in output_mode_list:
                 assert (