ref: improve code

api/src/backend/api/migrations/0010_scan_report_output.py

@@ -14,9 +14,4 @@ class Migration(migrations.Migration):
             name="output_path",
             field=models.CharField(blank=True, max_length=200, null=True),
         ),
-        migrations.AddField(
-            model_name="scan",
-            name="upload_to_s3",
-            field=models.BooleanField(blank=True, null=True),
-        ),
     ]

api/src/backend/api/models.py

@@ -415,7 +415,6 @@ class Scan(RowLevelSecurityProtectedModel):
         PeriodicTask, on_delete=models.CASCADE, null=True, blank=True
     )
     output_path = models.CharField(blank=True, null=True, max_length=200)
-    upload_to_s3 = models.BooleanField(blank=True, null=True)
     # TODO: mutelist foreign key
 
     class Meta(RowLevelSecurityProtectedModel.Meta):

api/src/backend/api/v1/views.py

@@ -87,6 +87,7 @@ from api.models import (
 from api.pagination import ComplianceOverviewPagination
 from api.rbac.permissions import Permissions, get_providers, get_role
 from api.rls import Tenant
+from tasks.jobs.export import get_s3_client
 from api.utils import validate_invitation
 from api.uuid_utils import datetime_to_uuid7
 from api.v1.serializers import (
@@ -1106,7 +1107,7 @@ class ScanViewSet(BaseRLSViewSet):
         request=ScanReportSerializer,
         responses={
             200: OpenApiResponse(description="Report obtained successfully"),
-            423: OpenApiResponse(
+            403: OpenApiResponse(
                 description="There is a problem with the AWS credentials"
             ),
         },
@@ -1121,28 +1122,14 @@ class ScanViewSet(BaseRLSViewSet):
                 {"detail": "No files found"}, status=status.HTTP_404_NOT_FOUND
             )
 
-        if scan_instance.upload_to_s3:
-            s3_client = None
+        if scan_instance.output_path.startswith("s3://"):
             try:
-                s3_client = boto3.client(
-                    "s3",
-                    aws_access_key_id=env.str("DJANGO_ARTIFACTS_AWS_ACCESS_KEY_ID"),
-                    aws_secret_access_key=env.str(
-                        "DJANGO_ARTIFACTS_AWS_SECRET_ACCESS_KEY"
-                    ),
-                    aws_session_token=env.str("DJANGO_ARTIFACTS_AWS_SESSION_TOKEN"),
-                    region_name=env.str("DJANGO_ARTIFACTS_AWS_DEFAULT_REGION"),
-                )
-                s3_client.list_buckets()
+                s3_client = get_s3_client()
             except (ClientError, NoCredentialsError, ParamValidationError):
-                try:
-                    s3_client = boto3.client("s3")
-                    s3_client.list_buckets()
-                except (ClientError, NoCredentialsError, ParamValidationError):
-                    return Response(
-                        {"detail": "There is a problem with the AWS credentials."},
-                        status=status.HTTP_423_LOCKED,
-                    )
+                return Response(
+                    {"detail": "There is a problem with the AWS credentials."},
+                    status=status.HTTP_403_FORBIDDEN,
+                )
 
             bucket_name = env.str("DJANGO_ARTIFACTS_AWS_S3_OUTPUT_BUCKET")
 

api/src/backend/config/django/base.py

@@ -208,3 +208,14 @@ CACHE_STALE_WHILE_REVALIDATE = env.int("DJANGO_STALE_WHILE_REVALIDATE", 60)
 
 TESTING = False
 FINDINGS_MAX_DAYS_IN_RANGE = env.int("DJANGO_FINDINGS_MAX_DAYS_IN_RANGE", 7)
+
+
+# API export settings
+TMP_OUTPUT_DIRECTORY = "/tmp/prowler_api_output"
+FINDINGS_BATCH_SIZE = 1000
+
+DJANGO_ARTIFACTS_AWS_S3_OUTPUT_BUCKET = env.str("DJANGO_ARTIFACTS_AWS_S3_OUTPUT_BUCKET", "")
+DJANGO_ARTIFACTS_AWS_ACCESS_KEY_ID = env.str("DJANGO_ARTIFACTS_AWS_ACCESS_KEY_ID", "")
+DJANGO_ARTIFACTS_AWS_SECRET_ACCESS_KEY = env.str("DJANGO_ARTIFACTS_AWS_SECRET_ACCESS_KEY", "")
+DJANGO_ARTIFACTS_AWS_SESSION_TOKEN = env.str("DJANGO_ARTIFACTS_AWS_SESSION_TOKEN", "")
+DJANGO_ARTIFACTS_AWS_DEFAULT_REGION = env.str("DJANGO_ARTIFACTS_AWS_DEFAULT_REGION", "")

api/src/backend/tasks/jobs/export.py

@@ -2,15 +2,16 @@ import os
 import zipfile
 
 import boto3
-from botocore.exceptions import ClientError
+import config.django.base as base
+from botocore.exceptions import ClientError, NoCredentialsError, ParamValidationError
 from celery.utils.log import get_task_logger
 from config.env import env
+
 from prowler.config.config import (
     csv_file_suffix,
     html_file_suffix,
     json_ocsf_file_suffix,
     output_file_timestamp,
-    tmp_output_directory,
 )
 from prowler.lib.outputs.csv.csv import CSV
 from prowler.lib.outputs.html.html import HTML
@@ -53,6 +54,41 @@ def _compress_output_files(output_directory: str) -> str:
     return zip_path
 
 
+def get_s3_client():
+    """
+    Create and return a boto3 S3 client using AWS credentials from environment variables.
+
+    This function attempts to initialize an S3 client by reading the AWS access key, secret key,
+    session token, and region from environment variables. It then validates the client by listing
+    available S3 buckets. If an error occurs during this process (for example, due to missing or
+    invalid credentials), it falls back to creating an S3 client without explicitly provided credentials,
+    which may rely on other configuration sources (e.g., IAM roles).
+
+    Returns:
+        boto3.client: A configured S3 client instance.
+
+    Raises:
+        ClientError, NoCredentialsError, or ParamValidationError if both attempts to create a client fail.
+    """
+    s3_client = None
+    try:
+        s3_client = boto3.client(
+            "s3",
+            aws_access_key_id=env.str("DJANGO_ARTIFACTS_AWS_ACCESS_KEY_ID"),
+            aws_secret_access_key=env.str(
+                "DJANGO_ARTIFACTS_AWS_SECRET_ACCESS_KEY"
+            ),
+            aws_session_token=env.str("DJANGO_ARTIFACTS_AWS_SESSION_TOKEN"),
+            region_name=env.str("DJANGO_ARTIFACTS_AWS_DEFAULT_REGION"),
+        )
+        s3_client.list_buckets()
+    except (ClientError, NoCredentialsError, ParamValidationError):
+        s3_client = boto3.client("s3")
+        s3_client.list_buckets()
+
+    return s3_client
+
+
 def _upload_to_s3(tenant_id: str, zip_path: str, scan_id: str) -> str:
     """
     Upload the specified ZIP file to an S3 bucket.
@@ -68,35 +104,25 @@ def _upload_to_s3(tenant_id: str, zip_path: str, scan_id: str) -> str:
     Raises:
         botocore.exceptions.ClientError: If the upload attempt to S3 fails for any reason.
     """
-    if not env.str("DJANGO_ARTIFACTS_AWS_S3_OUTPUT_BUCKET", ""):
+    if not base.DJANGO_ARTIFACTS_AWS_S3_OUTPUT_BUCKET:
         return
 
-    if env.str("DJANGO_ARTIFACTS_AWS_ACCESS_KEY_ID", ""):
-        s3 = boto3.client(
-            "s3",
-            aws_access_key_id=env.str("DJANGO_ARTIFACTS_AWS_ACCESS_KEY_ID"),
-            aws_secret_access_key=env.str("DJANGO_ARTIFACTS_AWS_SECRET_ACCESS_KEY"),
-            aws_session_token=env.str("DJANGO_ARTIFACTS_AWS_SESSION_TOKEN"),
-            region_name=env.str("DJANGO_ARTIFACTS_AWS_DEFAULT_REGION"),
-        )
-    else:
-        s3 = boto3.client("s3")
-
+    s3 = get_s3_client()
     s3_key = f"{tenant_id}/{scan_id}/{os.path.basename(zip_path)}"
     try:
         s3.upload_file(
             Filename=zip_path,
-            Bucket=env.str("DJANGO_ARTIFACTS_AWS_S3_OUTPUT_BUCKET"),
+            Bucket=base.DJANGO_ARTIFACTS_AWS_S3_OUTPUT_BUCKET,
             Key=s3_key,
         )
-        return f"s3://{env.str('DJANGO_ARTIFACTS_AWS_S3_OUTPUT_BUCKET')}/{s3_key}"
+        return f"s3://{base.DJANGO_ARTIFACTS_AWS_S3_OUTPUT_BUCKET}/{s3_key}"
     except ClientError as e:
         logger.error(f"S3 upload failed: {str(e)}")
         raise e
 
 
 def _generate_output_directory(
-    prowler_provider: object, tenant_id: str, scan_id: str
+    output_directory, prowler_provider: object, tenant_id: str, scan_id: str
 ) -> str:
     """
     Generate a file system path for the output directory of a prowler scan.
@@ -107,11 +133,11 @@ def _generate_output_directory(
     store the output files of a prowler scan.
 
     Note:
-        This function depends on two external variables:
-          - `tmp_output_directory`: The base directory where temporary outputs are stored.
+        This function depends on one external variable:
           - `output_file_timestamp`: A timestamp (as a string) used to uniquely identify the output.
 
     Args:
+        output_directory (str): The base output directory.
         prowler_provider (object): An identifier or descriptor for the prowler provider.
                                    Typically, this is a string indicating the provider (e.g., "aws").
         tenant_id (str): The unique identifier for the tenant.
@@ -121,10 +147,13 @@ def _generate_output_directory(
         str: The constructed file system path for the prowler scan output directory.
 
     Example:
-        >>> _generate_output_directory("aws", "tenant-1234", "scan-5678")
-        '/tmp/tenant-1234/scan-5678/prowler-output-aws-2023-02-15T12:34:56'
+        >>> _generate_output_directory("/tmp", "aws", "tenant-1234", "scan-5678")
+        '/tmp/tenant-1234/aws/scan-5678/prowler-output-2023-02-15T12:34:56'
     """
-    return (
-        f"{tmp_output_directory}/{tenant_id}/{scan_id}/prowler-output-"
+    path = (
+        f"{output_directory}/{tenant_id}/{scan_id}/prowler-output-"
         f"{prowler_provider}-{output_file_timestamp}"
     )
+    os.makedirs("/".join(path.split("/")[:-1]), exist_ok=True)
+
+    return path

api/src/backend/tasks/tasks.py

@@ -1,8 +1,7 @@
-import os
-
 from celery import chain, shared_task
+from celery.utils.log import get_task_logger
 from config.celery import RLSTask
-from api.utils import initialize_prowler_provider
+from config.django.base import FINDINGS_BATCH_SIZE, TMP_OUTPUT_DIRECTORY
 from django_celery_beat.models import PeriodicTask
 from tasks.jobs.connection import check_provider_connection
 from tasks.jobs.deletion import delete_provider, delete_tenant
@@ -18,8 +17,11 @@ from tasks.utils import get_next_execution_datetime
 from api.db_utils import rls_transaction
 from api.decorators import set_tenant
 from api.models import Finding, Provider, Scan, ScanSummary, StateChoices
+from api.utils import initialize_prowler_provider
 from prowler.lib.outputs.finding import Finding as FindingOutput
 
+logger = get_task_logger(__name__)
+
 
 @shared_task(base=RLSTask, name="provider-connection-check")
 @set_tenant
@@ -212,9 +214,14 @@ def generate_outputs(scan_id: str, provider_id: str, tenant_id: str):
         scan_id (str): The scan identifier.
         provider_id (str): The provider_id id to be used in generating outputs.
     """
+    # Initialize the prowler provider
+    prowler_provider = initialize_prowler_provider(Provider.objects.get(id=provider_id))
+
+    # Get the provider UID
+    provider_uid = Provider.objects.get(id=provider_id).uid
+
     # Generate and ensure the output directory exists
-    output_directory = _generate_output_directory(provider_id, tenant_id, scan_id)
-    os.makedirs("/".join(output_directory.split("/")[:-1]), exist_ok=True)
+    output_directory = _generate_output_directory(TMP_OUTPUT_DIRECTORY, provider_uid, tenant_id, scan_id)
 
     # Define auxiliary variables
     output_writers = {}
@@ -226,7 +233,7 @@ def generate_outputs(scan_id: str, provider_id: str, tenant_id: str):
     findings_qs = Finding.objects.filter(scan_id=scan_id).order_by("uid")
 
     # Process findings in batches
-    for batch, is_last_batch in batched(findings_qs.iterator(), 50):
+    for batch, is_last_batch in batched(findings_qs.iterator(), FINDINGS_BATCH_SIZE):
         finding_outputs = [
             FindingOutput.transform_api_finding(finding) for finding in batch
         ]
@@ -235,9 +242,7 @@ def generate_outputs(scan_id: str, provider_id: str, tenant_id: str):
         for mode, config in OUTPUT_FORMATS_MAPPING.items():
             kwargs = dict(config.get("kwargs", {}))
             if mode == "html":
-                kwargs["provider"] = initialize_prowler_provider(
-                    Provider.objects.get(id=provider_id)
-                )
+                kwargs["provider"] = prowler_provider
                 kwargs["stats"] = scan_summary
 
             writer_class = config["class"]
@@ -256,6 +261,8 @@ def generate_outputs(scan_id: str, provider_id: str, tenant_id: str):
 
             # Write the current batch using the writer
             writer.batch_write_data_to_file(**kwargs)
+
+            # TODO: Refactor the output classes to avoid this manual reset
             writer._data = []
 
     # Compress output files
@@ -270,14 +277,13 @@ def generate_outputs(scan_id: str, provider_id: str, tenant_id: str):
     else:
         uploaded = False
 
-    # Update the scan instance with the output path and upload status
-    Scan.objects.filter(id=scan_id).update(
-        output_path=output_directory, upload_to_s3=uploaded
-    )
+    # Update the scan instance with the output path
+    Scan.objects.filter(id=scan_id).update(output_path=output_directory)
+
+    logger.info(f"Scan output files generated, output location: {output_directory}")
 
     return {
-        "output_path": output_directory,
-        "upload_to_s3": uploaded,
+        "upload": uploaded,
         "scan_id": scan_id,
         "provider_id": provider_id,
     }

docker-compose.yml

@@ -7,6 +7,8 @@ services:
         required: false
     ports:
       - "${DJANGO_PORT:-8080}:${DJANGO_PORT:-8080}"
+    volumes:
+      - "/tmp/prowler_api_output:/tmp/prowler_api_output"
     depends_on:
       postgres:
         condition: service_healthy
@@ -65,6 +67,8 @@ services:
     env_file:
       - path: .env
         required: false
+    volumes:
+      - "/tmp/prowler_api_output:/tmp/prowler_api_output"
     depends_on:
       valkey:
         condition: service_healthy

prowler/config/config.py

@@ -60,7 +60,6 @@ aws_services_json_file = "aws_regions_by_service.json"
 
 default_output_directory = getcwd() + "/output"
 output_file_timestamp = timestamp.strftime("%Y%m%d%H%M%S")
-tmp_output_directory = "/tmp/prowler_api_output"
 timestamp_iso = timestamp.isoformat(sep=" ", timespec="seconds")
 csv_file_suffix = ".csv"
 json_file_suffix = ".json"

prowler/lib/outputs/finding.py

@@ -337,14 +337,15 @@ class Finding(BaseModel):
         output_data["uid"] = finding.uid
         output_data["status"] = Status(finding.status)
         output_data["status_extended"] = finding.status_extended
-        output_data["resource_uid"] = finding.resources.first().uid
-        output_data["resource_name"] = finding.resources.first().name
+        resource = finding.resources.first()
+        output_data["resource_uid"] = resource.uid
+        output_data["resource_name"] = resource.name
         output_data["resource_details"] = ""
-        resource_tags = finding.resources.first().tags.all()
+        resource_tags = resource.tags.all()
         output_data["resource_tags"] = unroll_tags(
             [{"key": tag.key, "value": tag.value} for tag in resource_tags]
         )
-        output_data["region"] = finding.resources.first().region
+        output_data["region"] = resource.region
         output_data["compliance"] = {}
 
         return cls(**output_data)