feat(threatscore): add UI + API changes

api/src/backend/api/v1/serializers.py

@@ -983,6 +983,14 @@ class ScanComplianceReportSerializer(serializers.Serializer):
         fields = ["id", "name"]
 
 
+class ScanThreatscoreReportSerializer(serializers.Serializer):
+    id = serializers.CharField(source="scan")
+
+    class Meta:
+        resource_name = "scan-reports"
+        fields = ["id"]
+
+
 class ResourceTagSerializer(RLSSerializer):
     """
     Serializer for the ResourceTag model
@@ -1808,14 +1816,6 @@ class ComplianceOverviewMetadataSerializer(serializers.Serializer):
         resource_name = "compliance-overviews-metadata"
 
 
-class ComplianceOverviewThreatscoreReportSerializer(serializers.Serializer):
-    id = serializers.CharField(source="scan")
-
-    class Meta:
-        resource_name = "scan-threatscore-reports"
-        fields = ["id"]
-
-
 # Overviews
 
 

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

@@ -137,7 +137,6 @@ from api.v1.serializers import (
     ComplianceOverviewDetailSerializer,
     ComplianceOverviewMetadataSerializer,
     ComplianceOverviewSerializer,
-    ComplianceOverviewThreatscoreReportSerializer,
     FindingDynamicFilterSerializer,
     FindingMetadataSerializer,
     FindingSerializer,
@@ -181,6 +180,7 @@ from api.v1.serializers import (
     ScanCreateSerializer,
     ScanReportSerializer,
     ScanSerializer,
+    ScanThreatscoreReportSerializer,
     ScanUpdateSerializer,
     ScheduleDailyCreateSerializer,
     TaskSerializer,
@@ -293,7 +293,7 @@ class SchemaView(SpectacularAPIView):
 
     def get(self, request, *args, **kwargs):
         spectacular_settings.TITLE = "Prowler API"
-        spectacular_settings.VERSION = "1.10.1"
+        spectacular_settings.VERSION = "1.10.2"
         spectacular_settings.DESCRIPTION = (
             "Prowler API specification.\n\nThis file is auto-generated."
         )
@@ -1497,6 +1497,22 @@ class ProviderViewSet(BaseRLSViewSet):
         },
         request=None,
     ),
+    threatscore=extend_schema(
+        tags=["Scan"],
+        summary="Retrieve threatscore report",
+        description="Download a specific threatscore report (e.g., 'prowler_threatscore_aws') as a PDF file.",
+        request=ScanThreatscoreReportSerializer,
+        responses={
+            200: OpenApiResponse(
+                description="PDF file containing the threatscore report"
+            ),
+            202: OpenApiResponse(description="The task is in progress"),
+            403: OpenApiResponse(description="There is a problem with credentials"),
+            404: OpenApiResponse(
+                description="The scan has no threatscore reports, or the threatscore report generation task has not started yet"
+            ),
+        },
+    ),
 )
 @method_decorator(CACHE_DECORATOR, name="list")
 @method_decorator(CACHE_DECORATOR, name="retrieve")
@@ -1553,6 +1569,10 @@ class ScanViewSet(BaseRLSViewSet):
             if hasattr(self, "response_serializer_class"):
                 return self.response_serializer_class
             return ScanComplianceReportSerializer
+        elif self.action == "threatscore":
+            if hasattr(self, "response_serializer_class"):
+                return self.response_serializer_class
+            return ScanThreatscoreReportSerializer
         return super().get_serializer_class()
 
     def partial_update(self, request, *args, **kwargs):
@@ -1784,6 +1804,45 @@ class ScanViewSet(BaseRLSViewSet):
         content, filename = loader
         return self._serve_file(content, filename, "text/csv")
 
+    @action(
+        detail=True,
+        methods=["get"],
+        url_name="threatscore",
+    )
+    def threatscore(self, request, pk=None):
+        scan = self.get_object()
+        running_resp = self._get_task_status(scan)
+        if running_resp:
+            return running_resp
+
+        if not scan.output_location:
+            return Response(
+                {
+                    "detail": "The scan has no reports, or the threatscore report generation task has not started yet."
+                },
+                status=status.HTTP_404_NOT_FOUND,
+            )
+
+        if scan.output_location.startswith("s3://"):
+            bucket = env.str("DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET", "")
+            key_prefix = scan.output_location.removeprefix(f"s3://{bucket}/")
+            prefix = os.path.join(
+                os.path.dirname(key_prefix),
+                "threatscore",
+                "*_threatscore_report.pdf",
+            )
+            loader = self._load_file(prefix, s3=True, bucket=bucket, list_objects=True)
+        else:
+            base = os.path.dirname(scan.output_location)
+            pattern = os.path.join(base, "threatscore", "*_threatscore_report.pdf")
+            loader = self._load_file(pattern, s3=False)
+
+        if isinstance(loader, Response):
+            return loader
+
+        content, filename = loader
+        return self._serve_file(content, filename, "application/pdf")
+
     def create(self, request, *args, **kwargs):
         input_serializer = self.get_serializer(data=request.data)
         input_serializer.is_valid(raise_exception=True)
@@ -3118,8 +3177,6 @@ class ComplianceOverviewViewSet(BaseRLSViewSet, TaskManagementMixin):
             return ComplianceOverviewAttributesSerializer
         elif self.action == "requirements":
             return ComplianceOverviewDetailSerializer
-        elif self.action == "threatscore":
-            return ComplianceOverviewThreatscoreReportSerializer
         return super().get_serializer_class()
 
     @extend_schema(exclude=True)
@@ -3427,75 +3484,6 @@ class ComplianceOverviewViewSet(BaseRLSViewSet, TaskManagementMixin):
         serializer = self.get_serializer(attribute_data, many=True)
         return Response(serializer.data, status=status.HTTP_200_OK)
 
-    @action(
-        detail=True,
-        methods=["get"],
-        url_path="threatscore",
-        url_name="threatscore",
-    )
-    def threatscore(self, request):
-        scan_id = request.query_params.get("filter[scan_id]")
-        if not scan_id:
-            raise ValidationError(
-                [
-                    {
-                        "detail": "This query parameter is required.",
-                        "status": 400,
-                        "source": {"pointer": "filter[scan_id]"},
-                        "code": "required",
-                    }
-                ]
-            )
-        try:
-            if task := self.get_task_response_if_running(
-                task_name="scan-threatscore-report",
-                task_kwargs={"tenant_id": self.request.tenant_id, "scan_id": scan_id},
-                raise_on_not_found=False,
-            ):
-                return task
-        except TaskFailedException:
-            return Response(
-                {"detail": "Task failed to generate threatscore report."},
-                status=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            )
-        try:
-            scan = Scan.objects.get(id=scan_id)
-        except Scan.DoesNotExist:
-            return Response(
-                {"detail": "Scan not found."},
-                status=status.HTTP_404_NOT_FOUND,
-            )
-
-        if not scan.output_location:
-            return Response(
-                {
-                    "detail": "The scan has no reports, or the report generation task has not started yet."
-                },
-                status=status.HTTP_404_NOT_FOUND,
-            )
-
-        if scan.output_location.startswith("s3://"):
-            bucket = env.str("DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET", "")
-            key_prefix = scan.output_location.removeprefix(f"s3://{bucket}/")
-            prefix = os.path.join(
-                os.path.dirname(key_prefix),
-                "compliance",
-                f"{scan.provider.provider}_threatscore.pdf",
-            )
-            loader = self._load_file(prefix, s3=True, bucket=bucket, list_objects=True)
-        else:
-            base = os.path.dirname(scan.output_location)
-            pattern = os.path.join(
-                base, "compliance", f"*_{scan.provider.provider}_threatscore.pdf"
-            )
-            loader = self._load_file(pattern, s3=False)
-
-        if isinstance(loader, Response):
-            return loader
-
-        content, filename = loader
-        return self._serve_file(content, filename, "application/pdf")
-
 
 @extend_schema(tags=["Overview"])
 @extend_schema_view(

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

@@ -240,6 +240,7 @@ def _generate_output_directory(
         >>> _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'
     """
     # Sanitize the prowler provider name to ensure it is a valid directory name
     prowler_provider_sanitized = re.sub(r"[^\w\-]", "-", prowler_provider)
@@ -256,4 +257,10 @@ def _generate_output_directory(
     )
     os.makedirs("/".join(compliance_path.split("/")[:-1]), exist_ok=True)
 
-    return path, compliance_path
+    threatscore_path = (
+        f"{output_directory}/{tenant_id}/{scan_id}/threatscore/prowler-output-"
+        f"{prowler_provider_sanitized}-{output_file_timestamp}"
+    )
+    os.makedirs("/".join(threatscore_path.split("/")[:-1]), exist_ok=True)
+
+    return path, compliance_path, threatscore_path

api/src/backend/tasks/tasks.py

@@ -281,7 +281,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
     )
 

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

@@ -420,8 +420,13 @@ 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")
     def test_scan_complete_tasks(
-        self, mock_outputs_task, mock_scan_summary_task, mock_compliance_tasks
+        self,
+        mock_threatscore_task,
+        mock_outputs_task,
+        mock_scan_summary_task,
+        mock_compliance_tasks,
     ):
         _perform_scan_complete_tasks("tenant-id", "scan-id", "provider-id")
         mock_compliance_tasks.assert_called_once_with(
@@ -436,3 +441,8 @@ class TestScanCompleteTasks:
             provider_id="provider-id",
             tenant_id="tenant-id",
         )
+        mock_threatscore_task.assert_called_once_with(
+            scan_id="scan-id",
+            provider_id="provider-id",
+            tenant_id="tenant-id",
+        )

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

@@ -0,0 +1,351 @@
+import uuid
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+from tasks.jobs.report import (
+    generate_threatscore_report,
+    generate_threatscore_report_job,
+)
+from tasks.tasks import generate_threatscore_report_task
+
+
+@pytest.mark.django_db
+class TestGenerateThreatscoreReport:
+    def setup_method(self):
+        self.scan_id = str(uuid.uuid4())
+        self.provider_id = str(uuid.uuid4())
+        self.tenant_id = str(uuid.uuid4())
+
+    def test_no_findings_returns_early(self):
+        with patch("tasks.jobs.report.ScanSummary.objects.filter") as mock_filter:
+            mock_filter.return_value.exists.return_value = False
+
+            result = generate_threatscore_report_job(
+                tenant_id=self.tenant_id,
+                scan_id=self.scan_id,
+                provider_id=self.provider_id,
+            )
+
+            assert result == {"upload": False}
+            mock_filter.assert_called_once_with(scan_id=self.scan_id)
+
+    @patch("tasks.jobs.report.rmtree")
+    @patch("tasks.jobs.report._upload_to_s3")
+    @patch("tasks.jobs.report.generate_threatscore_report")
+    @patch("tasks.jobs.report._generate_output_directory")
+    @patch("tasks.jobs.report.Provider.objects.get")
+    @patch("tasks.jobs.report.ScanSummary.objects.filter")
+    @patch("tasks.jobs.report.Scan.all_objects.filter")
+    def test_generate_threatscore_report_happy_path(
+        self,
+        mock_scan_update,
+        mock_scan_summary_filter,
+        mock_provider_get,
+        mock_generate_output_directory,
+        mock_generate_report,
+        mock_upload,
+        mock_rmtree,
+    ):
+        mock_scan_summary_filter.return_value.exists.return_value = True
+
+        mock_provider = MagicMock()
+        mock_provider.uid = "provider-uid"
+        mock_provider.provider = "aws"
+        mock_provider_get.return_value = mock_provider
+
+        mock_generate_output_directory.return_value = (
+            "/tmp/output",
+            "/tmp/compressed",
+            "/tmp/threatscore_path",
+        )
+
+        mock_upload.return_value = "s3://bucket/threatscore_report.pdf"
+
+        result = generate_threatscore_report_job(
+            tenant_id=self.tenant_id,
+            scan_id=self.scan_id,
+            provider_id=self.provider_id,
+        )
+
+        assert result == {"upload": True}
+        mock_generate_report.assert_called_once_with(
+            scan_id=self.scan_id,
+            compliance_id="prowler_threatscore_aws",
+            output_path="/tmp/threatscore_path_threatscore_report.pdf",
+            provider_id=self.provider_id,
+            only_failed=True,
+            min_risk_level=4,
+        )
+        mock_scan_update.return_value.update.assert_called_once_with(
+            output_location="s3://bucket/threatscore_report.pdf"
+        )
+        mock_rmtree.assert_called_once_with(
+            Path("/tmp/threatscore_path_threatscore_report.pdf").parent,
+            ignore_errors=True,
+        )
+
+    def test_generate_threatscore_report_fails_upload(self):
+        with (
+            patch("tasks.jobs.report.ScanSummary.objects.filter") as mock_filter,
+            patch("tasks.jobs.report.Provider.objects.get"),
+            patch("tasks.jobs.report._generate_output_directory") as mock_gen_dir,
+            patch("tasks.jobs.report.generate_threatscore_report"),
+            patch("tasks.jobs.report._upload_to_s3", return_value=None),
+            patch("tasks.jobs.report.Scan.all_objects.filter") as mock_scan_update,
+        ):
+            mock_filter.return_value.exists.return_value = True
+            mock_gen_dir.return_value = (
+                "/tmp/output",
+                "/tmp/compressed",
+                "/tmp/threatscore_path",
+            )
+
+            result = generate_threatscore_report_job(
+                tenant_id=self.tenant_id,
+                scan_id=self.scan_id,
+                provider_id=self.provider_id,
+            )
+
+            assert result == {"upload": False}
+            mock_scan_update.return_value.update.assert_called_once()
+
+    def test_generate_threatscore_report_logs_rmtree_exception(self, caplog):
+        with (
+            patch("tasks.jobs.report.ScanSummary.objects.filter") as mock_filter,
+            patch("tasks.jobs.report.Provider.objects.get"),
+            patch("tasks.jobs.report._generate_output_directory") as mock_gen_dir,
+            patch("tasks.jobs.report.generate_threatscore_report"),
+            patch(
+                "tasks.jobs.report._upload_to_s3", return_value="s3://bucket/report.pdf"
+            ),
+            patch("tasks.jobs.report.Scan.all_objects.filter"),
+            patch(
+                "tasks.jobs.report.rmtree", side_effect=Exception("Test deletion error")
+            ),
+        ):
+            mock_filter.return_value.exists.return_value = True
+            mock_gen_dir.return_value = (
+                "/tmp/output",
+                "/tmp/compressed",
+                "/tmp/threatscore_path",
+            )
+
+            with caplog.at_level("ERROR"):
+                generate_threatscore_report_job(
+                    tenant_id=self.tenant_id,
+                    scan_id=self.scan_id,
+                    provider_id=self.provider_id,
+                )
+                assert "Error deleting output files" in caplog.text
+
+    def test_generate_threatscore_report_azure_provider(self):
+        with (
+            patch("tasks.jobs.report.ScanSummary.objects.filter") as mock_filter,
+            patch("tasks.jobs.report.Provider.objects.get") as mock_provider_get,
+            patch("tasks.jobs.report._generate_output_directory") as mock_gen_dir,
+            patch("tasks.jobs.report.generate_threatscore_report") as mock_generate,
+            patch(
+                "tasks.jobs.report._upload_to_s3", return_value="s3://bucket/report.pdf"
+            ),
+            patch("tasks.jobs.report.Scan.all_objects.filter"),
+            patch("tasks.jobs.report.rmtree"),
+        ):
+            mock_filter.return_value.exists.return_value = True
+
+            mock_provider = MagicMock()
+            mock_provider.uid = "azure-provider-uid"
+            mock_provider.provider = "azure"
+            mock_provider_get.return_value = mock_provider
+
+            mock_gen_dir.return_value = (
+                "/tmp/output",
+                "/tmp/compressed",
+                "/tmp/threatscore_path",
+            )
+
+            generate_threatscore_report_job(
+                tenant_id=self.tenant_id,
+                scan_id=self.scan_id,
+                provider_id=self.provider_id,
+            )
+
+            mock_generate.assert_called_once_with(
+                scan_id=self.scan_id,
+                compliance_id="prowler_threatscore_azure",
+                output_path="/tmp/threatscore_path_threatscore_report.pdf",
+                provider_id=self.provider_id,
+                only_failed=True,
+                min_risk_level=4,
+            )
+
+
+@pytest.mark.django_db
+class TestGenerateThreatscoreReportFunction:
+    def setup_method(self):
+        self.scan_id = str(uuid.uuid4())
+        self.provider_id = str(uuid.uuid4())
+        self.compliance_id = "prowler_threatscore_aws"
+        self.output_path = "/tmp/test_threatscore_report.pdf"
+
+    @patch("tasks.jobs.report.initialize_prowler_provider")
+    @patch("tasks.jobs.report.Provider.objects.get")
+    @patch("tasks.jobs.report.Compliance.get_bulk")
+    @patch("tasks.jobs.report.Finding.all_objects.filter")
+    @patch("tasks.jobs.report.batched")
+    @patch("tasks.jobs.report.FindingOutput.transform_api_finding")
+    @patch("tasks.jobs.report.SimpleDocTemplate")
+    @patch("tasks.jobs.report.Image")
+    @patch("tasks.jobs.report.Spacer")
+    @patch("tasks.jobs.report.Paragraph")
+    @patch("tasks.jobs.report.PageBreak")
+    @patch("tasks.jobs.report.Table")
+    @patch("tasks.jobs.report.TableStyle")
+    @patch("tasks.jobs.report.plt.subplots")
+    @patch("tasks.jobs.report.io.BytesIO")
+    def test_generate_threatscore_report_success(
+        self,
+        mock_bytesio,
+        mock_subplots,
+        mock_table_style,
+        mock_table,
+        mock_page_break,
+        mock_paragraph,
+        mock_spacer,
+        mock_image,
+        mock_doc_template,
+        mock_transform_finding,
+        mock_batched,
+        mock_finding_filter,
+        mock_compliance_get_bulk,
+        mock_provider_get,
+        mock_initialize_provider,
+    ):
+        mock_provider = MagicMock()
+        mock_provider.provider = "aws"
+        mock_provider_get.return_value = mock_provider
+
+        prowler_provider = MagicMock()
+        mock_initialize_provider.return_value = prowler_provider
+
+        mock_compliance_obj = MagicMock()
+        mock_compliance_obj.Framework = "ProwlerThreatScore"
+        mock_compliance_obj.Version = "1.0"
+        mock_compliance_obj.Description = "Test Description"
+        mock_compliance_obj.Requirements = []
+        mock_compliance_get_bulk.return_value = {
+            self.compliance_id: mock_compliance_obj
+        }
+
+        mock_finding = MagicMock()
+        mock_finding.uid = "finding-1"
+        mock_finding_filter.return_value.order_by.return_value.iterator.return_value = [
+            mock_finding
+        ]
+
+        mock_batched.return_value = [([mock_finding], True)]
+
+        mock_transformed_finding = MagicMock()
+        mock_transformed_finding.check_id = "check-1"
+        mock_transformed_finding.status = "FAIL"
+        mock_transformed_finding.description = "Test finding"
+        mock_transformed_finding.metadata = MagicMock()
+        mock_transformed_finding.metadata.CheckTitle = "Test Check"
+        mock_transformed_finding.metadata.Severity = "HIGH"
+        mock_transformed_finding.resource_name = "test-resource"
+        mock_transformed_finding.region = "us-east-1"
+        mock_transform_finding.return_value = mock_transformed_finding
+
+        mock_doc = MagicMock()
+        mock_doc_template.return_value = mock_doc
+
+        mock_fig, mock_ax = MagicMock(), MagicMock()
+        mock_subplots.return_value = (mock_fig, mock_ax)
+        mock_buffer = MagicMock()
+        mock_bytesio.return_value = mock_buffer
+
+        mock_image.return_value = MagicMock()
+        mock_spacer.return_value = MagicMock()
+        mock_paragraph.return_value = MagicMock()
+        mock_page_break.return_value = MagicMock()
+        mock_table.return_value = MagicMock()
+        mock_table_style.return_value = MagicMock()
+
+        generate_threatscore_report(
+            scan_id=self.scan_id,
+            compliance_id=self.compliance_id,
+            output_path=self.output_path,
+            provider_id=self.provider_id,
+            only_failed=True,
+            min_risk_level=4,
+        )
+
+        mock_provider_get.assert_called_once_with(id=self.provider_id)
+        mock_initialize_provider.assert_called_once_with(mock_provider)
+        mock_compliance_get_bulk.assert_called_once_with("aws")
+        mock_finding_filter.assert_called_once_with(scan_id=self.scan_id)
+        mock_batched.assert_called_once()
+        mock_transform_finding.assert_called_once_with(mock_finding, prowler_provider)
+        mock_doc_template.assert_called_once()
+        mock_doc.build.assert_called_once()
+
+    @patch("tasks.jobs.report.initialize_prowler_provider")
+    @patch("tasks.jobs.report.Provider.objects.get")
+    @patch("tasks.jobs.report.Compliance.get_bulk")
+    @patch("tasks.jobs.report.Finding.all_objects.filter")
+    def test_generate_threatscore_report_exception_handling(
+        self,
+        mock_finding_filter,
+        mock_compliance_get_bulk,
+        mock_provider_get,
+        mock_initialize_provider,
+    ):
+        mock_provider_get.side_effect = Exception("Provider not found")
+
+        with pytest.raises(Exception, match="Provider not found"):
+            generate_threatscore_report(
+                scan_id=self.scan_id,
+                compliance_id=self.compliance_id,
+                output_path=self.output_path,
+                provider_id=self.provider_id,
+                only_failed=True,
+                min_risk_level=4,
+            )
+
+
+@pytest.mark.django_db
+class TestGenerateThreatscoreReportTask:
+    def setup_method(self):
+        self.scan_id = str(uuid.uuid4())
+        self.provider_id = str(uuid.uuid4())
+        self.tenant_id = str(uuid.uuid4())
+
+    @patch("tasks.tasks.generate_threatscore_report_job")
+    def test_generate_threatscore_report_task_calls_job(self, mock_generate_job):
+        mock_generate_job.return_value = {"upload": True}
+
+        result = generate_threatscore_report_task(
+            tenant_id=self.tenant_id,
+            scan_id=self.scan_id,
+            provider_id=self.provider_id,
+        )
+
+        assert result == {"upload": True}
+        mock_generate_job.assert_called_once_with(
+            tenant_id=self.tenant_id,
+            scan_id=self.scan_id,
+            provider_id=self.provider_id,
+        )
+
+    @patch("tasks.tasks.generate_threatscore_report_job")
+    def test_generate_threatscore_report_task_handles_job_exception(
+        self, mock_generate_job
+    ):
+        mock_generate_job.side_effect = Exception("Job failed")
+
+        with pytest.raises(Exception, match="Job failed"):
+            generate_threatscore_report_task(
+                tenant_id=self.tenant_id,
+                scan_id=self.scan_id,
+                provider_id=self.provider_id,
+            )