feat(export): add API export system (#6878)

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

@@ -0,0 +1,156 @@
+import os
+import zipfile
+
+import boto3
+import config.django.base as base
+from botocore.exceptions import ClientError, NoCredentialsError, ParamValidationError
+from celery.utils.log import get_task_logger
+from django.conf import settings
+
+from prowler.config.config import (
+    csv_file_suffix,
+    html_file_suffix,
+    json_ocsf_file_suffix,
+    output_file_timestamp,
+)
+from prowler.lib.outputs.csv.csv import CSV
+from prowler.lib.outputs.html.html import HTML
+from prowler.lib.outputs.ocsf.ocsf import OCSF
+
+logger = get_task_logger(__name__)
+
+
+# Predefined mapping for output formats and their configurations
+OUTPUT_FORMATS_MAPPING = {
+    "csv": {
+        "class": CSV,
+        "suffix": csv_file_suffix,
+        "kwargs": {},
+    },
+    "json-ocsf": {"class": OCSF, "suffix": json_ocsf_file_suffix, "kwargs": {}},
+    "html": {"class": HTML, "suffix": html_file_suffix, "kwargs": {"stats": {}}},
+}
+
+
+def _compress_output_files(output_directory: str) -> str:
+    """
+    Compress output files from all configured output formats into a ZIP archive.
+    Args:
+        output_directory (str): The directory where the output files are located.
+            The function looks up all known suffixes in OUTPUT_FORMATS_MAPPING
+            and compresses those files into a single ZIP.
+    Returns:
+        str: The full path to the newly created ZIP archive.
+    """
+    zip_path = f"{output_directory}.zip"
+
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for suffix in [config["suffix"] for config in OUTPUT_FORMATS_MAPPING.values()]:
+            zipf.write(
+                f"{output_directory}{suffix}",
+                f"output/{output_directory.split('/')[-1]}{suffix}",
+            )
+
+    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=settings.DJANGO_OUTPUT_S3_AWS_ACCESS_KEY_ID,
+            aws_secret_access_key=settings.DJANGO_OUTPUT_S3_AWS_SECRET_ACCESS_KEY,
+            aws_session_token=settings.DJANGO_OUTPUT_S3_AWS_SESSION_TOKEN,
+            region_name=settings.DJANGO_OUTPUT_S3_AWS_DEFAULT_REGION,
+        )
+        s3_client.list_buckets()
+    except (ClientError, NoCredentialsError, ParamValidationError, ValueError):
+        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.
+    If the S3 bucket environment variables are not configured,
+    the function returns None without performing an upload.
+    Args:
+        tenant_id (str): The tenant identifier, used as part of the S3 key prefix.
+        zip_path (str): The local file system path to the ZIP file to be uploaded.
+        scan_id (str): The scan identifier, used as part of the S3 key prefix.
+    Returns:
+        str: The S3 URI of the uploaded file (e.g., "s3://<bucket>/<key>") if successful.
+        None: If the required environment variables for the S3 bucket are not set.
+    Raises:
+        botocore.exceptions.ClientError: If the upload attempt to S3 fails for any reason.
+    """
+    if not base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET:
+        return
+
+    try:
+        s3 = get_s3_client()
+        s3_key = f"{tenant_id}/{scan_id}/{os.path.basename(zip_path)}"
+        s3.upload_file(
+            Filename=zip_path,
+            Bucket=base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET,
+            Key=s3_key,
+        )
+        return f"s3://{base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET}/{s3_key}"
+    except (ClientError, NoCredentialsError, ParamValidationError, ValueError) as e:
+        logger.error(f"S3 upload failed: {str(e)}")
+
+
+def _generate_output_directory(
+    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.
+
+    This function constructs the output directory path by combining a base
+    temporary output directory, the tenant ID, the scan ID, and details about
+    the prowler provider along with a timestamp. The resulting path is used to
+    store the output files of a prowler scan.
+
+    Note:
+        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.
+        scan_id (str): The unique identifier for the scan.
+
+    Returns:
+        str: The constructed file system path for the prowler scan output directory.
+
+    Example:
+        >>> _generate_output_directory("/tmp", "aws", "tenant-1234", "scan-5678")
+        '/tmp/tenant-1234/aws/scan-5678/prowler-output-2023-02-15T12:34:56'
+    """
+    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,14 +1,28 @@
-from celery import shared_task
+from shutil import rmtree
+
+from celery import chain, shared_task
+from celery.utils.log import get_task_logger
 from config.celery import RLSTask
+from config.django.base import DJANGO_FINDINGS_BATCH_SIZE, DJANGO_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
+from tasks.jobs.export import (
+    OUTPUT_FORMATS_MAPPING,
+    _compress_output_files,
+    _generate_output_directory,
+    _upload_to_s3,
+)
 from tasks.jobs.scan import aggregate_findings, perform_prowler_scan
-from tasks.utils import get_next_execution_datetime
+from tasks.utils import batched, get_next_execution_datetime
 
 from api.db_utils import rls_transaction
 from api.decorators import set_tenant
-from api.models import Scan, StateChoices
+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")
@@ -68,13 +82,20 @@ def perform_scan_task(
     Returns:
         dict: The result of the scan execution, typically including the status and results of the performed checks.
     """
-    return perform_prowler_scan(
+    result = perform_prowler_scan(
         tenant_id=tenant_id,
         scan_id=scan_id,
         provider_id=provider_id,
         checks_to_execute=checks_to_execute,
     )
 
+    chain(
+        perform_scan_summary_task.si(tenant_id, scan_id),
+        generate_outputs.si(scan_id, provider_id, tenant_id=tenant_id),
+    ).apply_async()
+
+    return result
+
 
 @shared_task(base=RLSTask, bind=True, name="scan-perform-scheduled", queue="scans")
 def perform_scheduled_scan_task(self, tenant_id: str, provider_id: str):
@@ -135,12 +156,11 @@ def perform_scheduled_scan_task(self, tenant_id: str, provider_id: str):
                 scheduler_task_id=periodic_task_instance.id,
             )
 
-    perform_scan_summary_task.apply_async(
-        kwargs={
-            "tenant_id": tenant_id,
-            "scan_id": str(scan_instance.id),
-        }
-    )
+    chain(
+        perform_scan_summary_task.si(tenant_id, scan_instance.id),
+        generate_outputs.si(str(scan_instance.id), provider_id, tenant_id=tenant_id),
+    ).apply_async()
+
     return result
 
 
@@ -152,3 +172,108 @@ def perform_scan_summary_task(tenant_id: str, scan_id: str):
 @shared_task(name="tenant-deletion", queue="deletion")
 def delete_tenant_task(tenant_id: str):
     return delete_tenant(pk=tenant_id)
+
+
+@shared_task(
+    base=RLSTask,
+    name="scan-report",
+    queue="scan-reports",
+)
+@set_tenant(keep_tenant=True)
+def generate_outputs(scan_id: str, provider_id: str, tenant_id: str):
+    """
+    Process findings in batches and generate output files in multiple formats.
+
+    This function retrieves findings associated with a scan, processes them
+    in batches of 50, and writes each batch to the corresponding output files.
+    It reuses output writer instances across batches, updates them with each
+    batch of transformed findings, and uses a flag to indicate when the final
+    batch is being processed. Finally, the output files are compressed and
+    uploaded to S3.
+
+    Args:
+        tenant_id (str): The tenant identifier.
+        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(
+        DJANGO_TMP_OUTPUT_DIRECTORY, provider_uid, tenant_id, scan_id
+    )
+
+    # Define auxiliary variables
+    output_writers = {}
+    scan_summary = FindingOutput._transform_findings_stats(
+        ScanSummary.objects.filter(scan_id=scan_id)
+    )
+
+    # Retrieve findings queryset
+    findings_qs = Finding.all_objects.filter(scan_id=scan_id).order_by("uid")
+
+    # Process findings in batches
+    for batch, is_last_batch in batched(
+        findings_qs.iterator(), DJANGO_FINDINGS_BATCH_SIZE
+    ):
+        finding_outputs = [
+            FindingOutput.transform_api_finding(finding, prowler_provider)
+            for finding in batch
+        ]
+
+        # Generate output files
+        for mode, config in OUTPUT_FORMATS_MAPPING.items():
+            kwargs = dict(config.get("kwargs", {}))
+            if mode == "html":
+                kwargs["provider"] = prowler_provider
+                kwargs["stats"] = scan_summary
+
+            writer_class = config["class"]
+            if writer_class in output_writers:
+                writer = output_writers[writer_class]
+                writer.transform(finding_outputs)
+                writer.close_file = is_last_batch
+            else:
+                writer = writer_class(
+                    findings=finding_outputs,
+                    file_path=output_directory,
+                    file_extension=config["suffix"],
+                    from_cli=False,
+                )
+                writer.close_file = is_last_batch
+                output_writers[writer_class] = writer
+
+            # 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
+    output_directory = _compress_output_files(output_directory)
+
+    # Save to configured storage
+    uploaded = _upload_to_s3(tenant_id, output_directory, scan_id)
+
+    if uploaded:
+        output_directory = uploaded
+        uploaded = True
+        # Remove the local files after upload
+        rmtree(DJANGO_TMP_OUTPUT_DIRECTORY, ignore_errors=True)
+    else:
+        uploaded = False
+
+    # Update the scan instance with the output path
+    Scan.all_objects.filter(id=scan_id).update(output_location=output_directory)
+
+    logger.info(f"Scan output files generated, output location: {output_directory}")
+
+    return {
+        "upload": uploaded,
+        "scan_id": scan_id,
+        "provider_id": provider_id,
+    }

api/src/backend/tasks/tests/test_utils.py

@@ -4,7 +4,7 @@ from unittest.mock import patch
 import pytest
 from django_celery_beat.models import IntervalSchedule, PeriodicTask
 from django_celery_results.models import TaskResult
-from tasks.utils import get_next_execution_datetime
+from tasks.utils import batched, get_next_execution_datetime
 
 
 @pytest.mark.django_db
@@ -74,3 +74,29 @@ class TestGetNextExecutionDatetime:
             get_next_execution_datetime(
                 task_id=task_result.task_id, provider_id="nonexistent"
             )
+
+
+class TestBatchedFunction:
+    def test_empty_iterable(self):
+        result = list(batched([], 3))
+        assert result == [([], True)]
+
+    def test_exact_batches(self):
+        result = list(batched([1, 2, 3, 4], 2))
+        expected = [([1, 2], False), ([3, 4], False), ([], True)]
+        assert result == expected
+
+    def test_inexact_batches(self):
+        result = list(batched([1, 2, 3, 4, 5], 2))
+        expected = [([1, 2], False), ([3, 4], False), ([5], True)]
+        assert result == expected
+
+    def test_batch_size_one(self):
+        result = list(batched([1, 2, 3], 1))
+        expected = [([1], False), ([2], False), ([3], False), ([], True)]
+        assert result == expected
+
+    def test_batch_size_greater_than_length(self):
+        result = list(batched([1, 2, 3], 5))
+        expected = [([1, 2, 3], True)]
+        assert result == expected

api/src/backend/tasks/utils.py

@@ -24,3 +24,27 @@ def get_next_execution_datetime(task_id: int, provider_id: str) -> datetime:
     )
 
     return current_scheduled_time + timedelta(**{interval.period: interval.every})
+
+
+def batched(iterable, batch_size):
+    """
+    Yield successive batches from an iterable.
+
+    Args:
+        iterable: An iterable source of items.
+        batch_size (int): The number of items per batch.
+
+    Yields:
+        tuple: A pair (batch, is_last_batch) where:
+            - batch (list): A list of items (with length equal to batch_size,
+              except possibly for the last batch).
+            - is_last_batch (bool): True if this is the final batch, False otherwise.
+    """
+    batch = []
+    for item in iterable:
+        batch.append(item)
+        if len(batch) == batch_size:
+            yield batch, False
+            batch = []
+
+    yield batch, True