feat(ens): support PDF reporting (#9158)

Co-authored-by: alejandrobailo <alejandrobailo94@gmail.com>

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

@@ -241,36 +241,33 @@ def _upload_to_s3(
         logger.error(f"S3 upload failed: {str(e)}")
 
 
-def _generate_output_directory(
-    output_directory, prowler_provider: object, tenant_id: str, scan_id: str
-) -> tuple[str, str, str]:
+def _build_output_path(
+    output_directory: str,
+    prowler_provider: str,
+    tenant_id: str,
+    scan_id: str,
+    subdirectory: str = None,
+) -> 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.
+    Build a file system path for the output directory of a prowler scan.
 
     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").
+        prowler_provider (str): 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.
+        subdirectory (str, optional): Optional subdirectory to include in the path
+                                     (e.g., "compliance", "threatscore", "ens").
 
     Returns:
-        str: The constructed file system path for the prowler scan output directory.
+        str: The constructed path with directory created.
 
     Example:
-        >>> _generate_output_directory("/tmp", "aws", "tenant-1234", "scan-5678")
-        '/tmp/tenant-1234/aws/scan-5678/prowler-output-2023-02-15T12:34:56',
-        '/tmp/tenant-1234/aws/scan-5678/compliance/prowler-output-2023-02-15T12:34:56'
-        '/tmp/tenant-1234/aws/scan-5678/threatscore/prowler-output-2023-02-15T12:34:56'
+        >>> _build_output_path("/tmp", "aws", "tenant-1234", "scan-5678")
+        '/tmp/tenant-1234/scan-5678/prowler-output-aws-20230215123456'
+        >>> _build_output_path("/tmp", "aws", "tenant-1234", "scan-5678", "threatscore")
+        '/tmp/tenant-1234/scan-5678/threatscore/prowler-output-aws-20230215123456'
     """
     # Sanitize the prowler provider name to ensure it is a valid directory name
     prowler_provider_sanitized = re.sub(r"[^\w\-]", "-", prowler_provider)
@@ -279,22 +276,102 @@ def _generate_output_directory(
         started_at = Scan.objects.get(id=scan_id).started_at
 
     timestamp = started_at.strftime("%Y%m%d%H%M%S")
-    path = (
-        f"{output_directory}/{tenant_id}/{scan_id}/prowler-output-"
-        f"{prowler_provider_sanitized}-{timestamp}"
-    )
+
+    if subdirectory:
+        path = (
+            f"{output_directory}/{tenant_id}/{scan_id}/{subdirectory}/prowler-output-"
+            f"{prowler_provider_sanitized}-{timestamp}"
+        )
+    else:
+        path = (
+            f"{output_directory}/{tenant_id}/{scan_id}/prowler-output-"
+            f"{prowler_provider_sanitized}-{timestamp}"
+        )
+
+    # Create directory for the path if it doesn't exist
     os.makedirs("/".join(path.split("/")[:-1]), exist_ok=True)
 
-    compliance_path = (
-        f"{output_directory}/{tenant_id}/{scan_id}/compliance/prowler-output-"
-        f"{prowler_provider_sanitized}-{timestamp}"
-    )
-    os.makedirs("/".join(compliance_path.split("/")[:-1]), exist_ok=True)
+    return path
 
-    threatscore_path = (
-        f"{output_directory}/{tenant_id}/{scan_id}/threatscore/prowler-output-"
-        f"{prowler_provider_sanitized}-{timestamp}"
-    )
-    os.makedirs("/".join(threatscore_path.split("/")[:-1]), exist_ok=True)
 
-    return path, compliance_path, threatscore_path
+def _generate_compliance_output_directory(
+    output_directory: str,
+    prowler_provider: str,
+    tenant_id: str,
+    scan_id: str,
+    compliance_framework: str,
+) -> str:
+    """
+    Generate a file system path for a compliance framework output directory.
+
+    This function constructs the output directory path specifically for a compliance
+    framework (e.g., "threatscore", "ens") by combining a base temporary output directory,
+    the tenant ID, the scan ID, the compliance framework name, and details about the
+    prowler provider along with a timestamp.
+
+    Args:
+        output_directory (str): The base output directory.
+        prowler_provider (str): 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.
+        compliance_framework (str): The compliance framework name (e.g., "threatscore", "ens").
+
+    Returns:
+        str: The path for the compliance framework output directory.
+
+    Example:
+        >>> _generate_compliance_output_directory("/tmp", "aws", "tenant-1234", "scan-5678", "threatscore")
+        '/tmp/tenant-1234/scan-5678/threatscore/prowler-output-aws-20230215123456'
+        >>> _generate_compliance_output_directory("/tmp", "aws", "tenant-1234", "scan-5678", "ens")
+        '/tmp/tenant-1234/scan-5678/ens/prowler-output-aws-20230215123456'
+    """
+    return _build_output_path(
+        output_directory,
+        prowler_provider,
+        tenant_id,
+        scan_id,
+        subdirectory=compliance_framework,
+    )
+
+
+def _generate_output_directory(
+    output_directory: str,
+    prowler_provider: str,
+    tenant_id: str,
+    scan_id: str,
+) -> tuple[str, str]:
+    """
+    Generate file system paths for the standard and compliance output directories of a prowler scan.
+
+    This function constructs both the standard output directory path and the compliance
+    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.
+
+    Args:
+        output_directory (str): The base output directory.
+        prowler_provider (str): 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:
+        tuple[str, str]: A tuple containing (standard_path, compliance_path).
+
+    Example:
+        >>> _generate_output_directory("/tmp", "aws", "tenant-1234", "scan-5678")
+        ('/tmp/tenant-1234/scan-5678/prowler-output-aws-20230215123456',
+         '/tmp/tenant-1234/scan-5678/compliance/prowler-output-aws-20230215123456')
+    """
+    standard_path = _build_output_path(
+        output_directory, prowler_provider, tenant_id, scan_id
+    )
+    compliance_path = _build_output_path(
+        output_directory,
+        prowler_provider,
+        tenant_id,
+        scan_id,
+        subdirectory="compliance",
+    )
+
+    return standard_path, compliance_path

api/src/backend/tasks/tasks.py

@@ -35,7 +35,7 @@ from tasks.jobs.lighthouse_providers import (
     refresh_lighthouse_provider_models,
 )
 from tasks.jobs.muting import mute_historical_findings
-from tasks.jobs.report import generate_threatscore_report_job
+from tasks.jobs.report import generate_compliance_reports_job
 from tasks.jobs.scan import (
     aggregate_findings,
     create_compliance_requirements,
@@ -75,7 +75,8 @@ def _perform_scan_complete_tasks(tenant_id: str, scan_id: str, provider_id: str)
             scan_id=scan_id, provider_id=provider_id, tenant_id=tenant_id
         ),
         group(
-            generate_threatscore_report_task.si(
+            # Use optimized task that generates both reports with shared queries
+            generate_compliance_reports_task.si(
                 tenant_id=tenant_id, scan_id=scan_id, provider_id=provider_id
             ),
             check_integrations_task.si(
@@ -319,7 +320,7 @@ def generate_outputs_task(scan_id: str, provider_id: str, tenant_id: str):
 
     frameworks_bulk = Compliance.get_bulk(provider_type)
     frameworks_avail = get_compliance_frameworks(provider_type)
-    out_dir, comp_dir, _ = _generate_output_directory(
+    out_dir, comp_dir = _generate_output_directory(
         DJANGO_TMP_OUTPUT_DIRECTORY, provider_uid, tenant_id, scan_id
     )
 
@@ -686,19 +687,33 @@ def jira_integration_task(
 
 @shared_task(
     base=RLSTask,
-    name="scan-threatscore-report",
+    name="scan-compliance-reports",
     queue="scan-reports",
 )
-def generate_threatscore_report_task(tenant_id: str, scan_id: str, provider_id: str):
+def generate_compliance_reports_task(tenant_id: str, scan_id: str, provider_id: str):
     """
-    Task to generate a threatscore report for a given scan.
+    Optimized task to generate both ThreatScore and ENS reports with shared queries.
+
+    This task is more efficient than running generate_threatscore_report_task and
+    generate_ens_report_task separately because it reuses database queries:
+    - Provider object fetched once (instead of twice)
+    - Requirement statistics aggregated once (instead of twice)
+    - Can reduce database load by up to 50%
+
     Args:
         tenant_id (str): The tenant identifier.
         scan_id (str): The scan identifier.
         provider_id (str): The provider identifier.
+
+    Returns:
+        dict: Results for both reports containing upload status and paths.
     """
-    return generate_threatscore_report_job(
-        tenant_id=tenant_id, scan_id=scan_id, provider_id=provider_id
+    return generate_compliance_reports_job(
+        tenant_id=tenant_id,
+        scan_id=scan_id,
+        provider_id=provider_id,
+        generate_threatscore=True,
+        generate_ens=True,
     )
 
 

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

@@ -9,6 +9,7 @@ import pytest
 from botocore.exceptions import ClientError
 from tasks.jobs.export import (
     _compress_output_files,
+    _generate_compliance_output_directory,
     _generate_output_directory,
     _upload_to_s3,
     get_s3_client,
@@ -168,17 +169,34 @@ class TestOutputs:
         provider = "aws"
         expected_timestamp = "20230615103045"
 
-        path, compliance, threatscore = _generate_output_directory(
+        # Test _generate_output_directory (returns standard and compliance paths)
+        path, compliance = _generate_output_directory(
             base_dir, provider, tenant_id, scan_id
         )
 
         assert os.path.isdir(os.path.dirname(path))
         assert os.path.isdir(os.path.dirname(compliance))
-        assert os.path.isdir(os.path.dirname(threatscore))
-
         assert path.endswith(f"{provider}-{expected_timestamp}")
         assert compliance.endswith(f"{provider}-{expected_timestamp}")
+        assert "/compliance/" in compliance
+
+        # Test _generate_compliance_output_directory with "threatscore"
+        threatscore = _generate_compliance_output_directory(
+            base_dir, provider, tenant_id, scan_id, compliance_framework="threatscore"
+        )
+
+        assert os.path.isdir(os.path.dirname(threatscore))
         assert threatscore.endswith(f"{provider}-{expected_timestamp}")
+        assert "/threatscore/" in threatscore
+
+        # Test _generate_compliance_output_directory with "ens"
+        ens = _generate_compliance_output_directory(
+            base_dir, provider, tenant_id, scan_id, compliance_framework="ens"
+        )
+
+        assert os.path.isdir(os.path.dirname(ens))
+        assert ens.endswith(f"{provider}-{expected_timestamp}")
+        assert "/ens/" in ens
 
     @patch("tasks.jobs.export.rls_transaction")
     @patch("tasks.jobs.export.Scan")
@@ -201,14 +219,25 @@ class TestOutputs:
         provider = "aws/test@check"
         expected_timestamp = "20230615103045"
 
-        path, compliance, threatscore = _generate_output_directory(
+        # Test provider name sanitization with _generate_output_directory
+        path, compliance = _generate_output_directory(
             base_dir, provider, tenant_id, scan_id
         )
 
         assert os.path.isdir(os.path.dirname(path))
         assert os.path.isdir(os.path.dirname(compliance))
-        assert os.path.isdir(os.path.dirname(threatscore))
-
         assert path.endswith(f"aws-test-check-{expected_timestamp}")
         assert compliance.endswith(f"aws-test-check-{expected_timestamp}")
+
+        # Test provider name sanitization with _generate_compliance_output_directory
+        threatscore = _generate_compliance_output_directory(
+            base_dir, provider, tenant_id, scan_id, compliance_framework="threatscore"
+        )
+        ens = _generate_compliance_output_directory(
+            base_dir, provider, tenant_id, scan_id, compliance_framework="ens"
+        )
+
+        assert os.path.isdir(os.path.dirname(threatscore))
+        assert os.path.isdir(os.path.dirname(ens))
         assert threatscore.endswith(f"aws-test-check-{expected_timestamp}")
+        assert ens.endswith(f"aws-test-check-{expected_timestamp}")

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

@@ -109,7 +109,6 @@ class TestGenerateOutputs:
                 return_value=(
                     "/tmp/test/out-dir",
                     "/tmp/test/comp-dir",
-                    "/tmp/test/threat-dir",
                 ),
             ),
             patch("tasks.tasks.Scan.all_objects.filter") as mock_scan_update,
@@ -139,7 +138,7 @@ class TestGenerateOutputs:
             patch("tasks.tasks.Finding.all_objects.filter") as mock_findings,
             patch(
                 "tasks.tasks._generate_output_directory",
-                return_value=("/tmp/test/out", "/tmp/test/comp", "/tmp/test/threat"),
+                return_value=("/tmp/test/out", "/tmp/test/comp"),
             ),
             patch("tasks.tasks.FindingOutput._transform_findings_stats"),
             patch("tasks.tasks.FindingOutput.transform_api_finding"),
@@ -209,7 +208,7 @@ class TestGenerateOutputs:
             patch("tasks.tasks.Finding.all_objects.filter") as mock_findings,
             patch(
                 "tasks.tasks._generate_output_directory",
-                return_value=("/tmp/test/out", "/tmp/test/comp", "/tmp/test/threat"),
+                return_value=("/tmp/test/out", "/tmp/test/comp"),
             ),
             patch(
                 "tasks.tasks.FindingOutput._transform_findings_stats",
@@ -289,7 +288,6 @@ class TestGenerateOutputs:
                 return_value=(
                     "/tmp/test/outdir",
                     "/tmp/test/compdir",
-                    "/tmp/test/threatdir",
                 ),
             ),
             patch("tasks.tasks._compress_output_files", return_value="outdir.zip"),
@@ -368,7 +366,6 @@ class TestGenerateOutputs:
                 return_value=(
                     "/tmp/test/outdir",
                     "/tmp/test/compdir",
-                    "/tmp/test/threatdir",
                 ),
             ),
             patch("tasks.tasks.FindingOutput._transform_findings_stats"),
@@ -436,7 +433,7 @@ class TestGenerateOutputs:
             patch("tasks.tasks.Finding.all_objects.filter") as mock_findings,
             patch(
                 "tasks.tasks._generate_output_directory",
-                return_value=("/tmp/test/out", "/tmp/test/comp", "/tmp/test/threat"),
+                return_value=("/tmp/test/out", "/tmp/test/comp"),
             ),
             patch(
                 "tasks.tasks.FindingOutput._transform_findings_stats",
@@ -494,7 +491,7 @@ class TestGenerateOutputs:
             patch("tasks.tasks.Finding.all_objects.filter") as mock_findings,
             patch(
                 "tasks.tasks._generate_output_directory",
-                return_value=("/tmp/test/out", "/tmp/test/comp", "/tmp/test/threat"),
+                return_value=("/tmp/test/out", "/tmp/test/comp"),
             ),
             patch("tasks.tasks.FindingOutput._transform_findings_stats"),
             patch("tasks.tasks.FindingOutput.transform_api_finding"),
@@ -535,34 +532,45 @@ class TestScanCompleteTasks:
     @patch("tasks.tasks.create_compliance_requirements_task.apply_async")
     @patch("tasks.tasks.perform_scan_summary_task.si")
     @patch("tasks.tasks.generate_outputs_task.si")
-    @patch("tasks.tasks.generate_threatscore_report_task.si")
+    @patch("tasks.tasks.generate_compliance_reports_task.si")
     @patch("tasks.tasks.check_integrations_task.si")
     def test_scan_complete_tasks(
         self,
         mock_check_integrations_task,
-        mock_threatscore_task,
+        mock_compliance_reports_task,
         mock_outputs_task,
         mock_scan_summary_task,
-        mock_compliance_tasks,
+        mock_compliance_requirements_task,
     ):
+        """Test that scan complete tasks are properly orchestrated with optimized reports."""
         _perform_scan_complete_tasks("tenant-id", "scan-id", "provider-id")
-        mock_compliance_tasks.assert_called_once_with(
+
+        # Verify compliance requirements task is called
+        mock_compliance_requirements_task.assert_called_once_with(
             kwargs={"tenant_id": "tenant-id", "scan_id": "scan-id"},
         )
+
+        # Verify scan summary task is called
         mock_scan_summary_task.assert_called_once_with(
             scan_id="scan-id",
             tenant_id="tenant-id",
         )
+
+        # Verify outputs task is called
         mock_outputs_task.assert_called_once_with(
             scan_id="scan-id",
             provider_id="provider-id",
             tenant_id="tenant-id",
         )
-        mock_threatscore_task.assert_called_once_with(
+
+        # Verify optimized compliance reports task is called (replaces individual tasks)
+        mock_compliance_reports_task.assert_called_once_with(
             tenant_id="tenant-id",
             scan_id="scan-id",
             provider_id="provider-id",
         )
+
+        # Verify integrations task is called
         mock_check_integrations_task.assert_called_once_with(
             tenant_id="tenant-id",
             provider_id="provider-id",
@@ -738,7 +746,7 @@ class TestCheckIntegrationsTask:
         mock_initialize_provider.return_value = MagicMock()
         mock_compliance_bulk.return_value = {}
         mock_get_frameworks.return_value = []
-        mock_generate_dir.return_value = ("out-dir", "comp-dir", "threat-dir")
+        mock_generate_dir.return_value = ("out-dir", "comp-dir")
         mock_transform_stats.return_value = {"stats": "data"}
 
         # Mock findings
@@ -863,7 +871,7 @@ class TestCheckIntegrationsTask:
         mock_initialize_provider.return_value = MagicMock()
         mock_compliance_bulk.return_value = {}
         mock_get_frameworks.return_value = []
-        mock_generate_dir.return_value = ("out-dir", "comp-dir", "threat-dir")
+        mock_generate_dir.return_value = ("out-dir", "comp-dir")
         mock_transform_stats.return_value = {"stats": "data"}
 
         # Mock findings
@@ -979,7 +987,7 @@ class TestCheckIntegrationsTask:
         mock_initialize_provider.return_value = MagicMock()
         mock_compliance_bulk.return_value = {}
         mock_get_frameworks.return_value = []
-        mock_generate_dir.return_value = ("out-dir", "comp-dir", "threat-dir")
+        mock_generate_dir.return_value = ("out-dir", "comp-dir")
         mock_transform_stats.return_value = {"stats": "data"}
 
         # Mock findings