feat(api): revert changes

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

@@ -1,8 +1,12 @@
 import glob
+import json
+import logging
 import os
+import tempfile
 from datetime import datetime, timedelta, timezone
 from urllib.parse import urljoin
 
+import boto3
 import sentry_sdk
 from allauth.socialaccount.models import SocialAccount, SocialApp
 from allauth.socialaccount.providers.github.views import GitHubOAuth2Adapter
@@ -48,7 +52,8 @@ from rest_framework.exceptions import (
     ValidationError,
 )
 from rest_framework.generics import GenericAPIView, get_object_or_404
-from rest_framework.permissions import SAFE_METHODS
+from rest_framework.permissions import SAFE_METHODS, IsAuthenticated
+from rest_framework.viewsets import ViewSet
 from rest_framework_json_api.views import RelationshipView, Response
 from rest_framework_simplejwt.exceptions import InvalidToken, TokenError
 from tasks.beat import schedule_provider_scan
@@ -59,6 +64,7 @@ from tasks.tasks import (
     check_provider_connection_task,
     delete_provider_task,
     delete_tenant_task,
+    generate_threatscore_report_task,
     perform_scan_task,
 )
 
@@ -1472,6 +1478,15 @@ class ProviderViewSet(BaseRLSViewSet):
         },
         request=None,
     ),
+    threatscore_report=extend_schema(
+        tags=["Scan"],
+        summary="Generate ThreatScore report",
+        description="Generate a ThreatScore report for a specific scan",
+        request=None,
+        responses={
+            202: OpenApiResponse(description="ThreatScore report generation started")
+        },
+    ),
 )
 @method_decorator(CACHE_DECORATOR, name="list")
 @method_decorator(CACHE_DECORATOR, name="retrieve")
@@ -1528,6 +1543,8 @@ class ScanViewSet(BaseRLSViewSet):
             if hasattr(self, "response_serializer_class"):
                 return self.response_serializer_class
             return ScanComplianceReportSerializer
+        elif self.action == "threatscore_report":
+            return None
         return super().get_serializer_class()
 
     def partial_update(self, request, *args, **kwargs):
@@ -1792,6 +1809,39 @@ class ScanViewSet(BaseRLSViewSet):
             },
         )
 
+    @action(detail=True, methods=["post"], url_path="threatscore-report")
+    def threatscore_report(self, request, pk=None):
+        scan = self.get_object()
+        tenant_id = self.request.tenant_id
+        provider_id = str(scan.provider_id)
+        scan_id = str(scan.id)
+        compliance_id = request.data.get("compliance_id")
+        output_path = request.data.get(
+            "output_path", f"/tmp/threatscore_report_{scan_id}.pdf"
+        )
+        only_failed = request.data.get("only_failed", True)
+        min_risk_level = request.data.get("min_risk_level", 4)
+        if not compliance_id:
+            return Response(
+                {"detail": "compliance_id is required"},
+                status=status.HTTP_400_BAD_REQUEST,
+            )
+        generate_threatscore_report_task.apply_async(
+            kwargs={
+                "tenant_id": tenant_id,
+                "scan_id": scan_id,
+                "provider_id": provider_id,
+                "compliance_id": compliance_id,
+                "output_path": output_path,
+                "only_failed": only_failed,
+                "min_risk_level": min_risk_level,
+            }
+        )
+        return Response(
+            {"detail": "ThreatScore report generation started."},
+            status=status.HTTP_202_ACCEPTED,
+        )
+
 
 @extend_schema_view(
     list=extend_schema(
@@ -3646,3 +3696,173 @@ class ProcessorViewSet(BaseRLSViewSet):
         elif self.action == "partial_update":
             return ProcessorUpdateSerializer
         return super().get_serializer_class()
+
+
+class ReportsViewSet(ViewSet):
+    permission_classes = [IsAuthenticated]
+    logger = logging.getLogger("prowler.reports")
+
+    def _get_task_and_validate(self, report_id, user):
+        try:
+            task = Task.objects.get(id=report_id)
+        except Task.DoesNotExist:
+            return None, Response(
+                {"detail": "Report not found"}, status=status.HTTP_404_NOT_FOUND
+            )
+
+        scan_id = None
+        try:
+            kwargs = json.loads(task.task_runner_task.task_kwargs.replace("'", '"'))
+            scan_id = kwargs.get("scan_id")
+        except Exception:
+            pass
+        if scan_id:
+            try:
+                scan = Scan.objects.get(id=scan_id)
+                if scan.tenant_id != user.tenant_id:
+                    return None, Response(
+                        {"detail": "Forbidden"}, status=status.HTTP_403_FORBIDDEN
+                    )
+            except Scan.DoesNotExist:
+                return None, Response(
+                    {"detail": "Scan not found"}, status=status.HTTP_404_NOT_FOUND
+                )
+        return task, None
+
+    @action(detail=False, methods=["post"], url_path="pdf/generate")
+    def pdf_generate(self, request):
+        scan_id = request.data.get("scan_id")
+        compliance_id = request.data.get("compliance_id")
+        options = request.data.get("options", {})
+        only_failed = options.get("only_failed", True)
+        min_risk_level = options.get("min_risk_level", 4)
+        if not scan_id or not compliance_id:
+            return Response(
+                {"detail": "scan_id and compliance_id are required"},
+                status=status.HTTP_400_BAD_REQUEST,
+            )
+        try:
+            scan = Scan.objects.get(id=scan_id)
+            if scan.tenant_id != request.user.tenant_id:
+                return Response(
+                    {"detail": "Forbidden"}, status=status.HTTP_403_FORBIDDEN
+                )
+        except Scan.DoesNotExist:
+            return Response(
+                {"detail": "Scan not found"}, status=status.HTTP_404_NOT_FOUND
+            )
+        tenant_id = scan.tenant_id
+        provider_id = scan.provider_id
+        output_path = f"/tmp/threatscore_report_{scan_id}.pdf"
+        celery_task = generate_threatscore_report_task.apply_async(
+            kwargs={
+                "tenant_id": tenant_id,
+                "scan_id": scan_id,
+                "provider_id": provider_id,
+                "compliance_id": compliance_id,
+                "output_path": output_path,
+                "only_failed": only_failed,
+                "min_risk_level": min_risk_level,
+            }
+        )
+        self.logger.info(
+            f"User {request.user.id} initiated ThreatScore report for scan {scan_id} (task {celery_task.id})"
+        )
+        return Response(
+            {"report_id": celery_task.id, "status": "initiated"},
+            status=status.HTTP_202_ACCEPTED,
+        )
+
+    @action(detail=False, methods=["get"], url_path="pdf/status/(?P<report_id>[^/]+)")
+    def pdf_status(self, request, report_id=None):
+        task, error = self._get_task_and_validate(report_id, request.user)
+        if error:
+            return error
+        state = task.state
+        output_path = None
+        try:
+            kwargs = json.loads(task.task_runner_task.task_kwargs.replace("'", '"'))
+            output_path = kwargs.get("output_path")
+        except Exception:
+            pass
+        if state == "SUCCESS":
+            status_str = "ready"
+        elif state in ("PENDING", "RECEIVED", "STARTED", "RETRY"):
+            status_str = "processing"
+        elif state == "FAILURE":
+            status_str = "failed"
+        else:
+            status_str = "initiated"
+        resp = {"status": status_str}
+        if status_str == "ready" and output_path:
+            resp["output_path"] = output_path
+        return Response(resp, status=status.HTTP_200_OK)
+
+    @action(detail=False, methods=["get"], url_path="pdf/download/(?P<report_id>[^/]+)")
+    def pdf_download(self, request, report_id=None):
+        task, error = self._get_task_and_validate(report_id, request.user)
+        if error:
+            return error
+        if task.state != "SUCCESS":
+            return Response(
+                {"detail": "Report not ready"}, status=status.HTTP_202_ACCEPTED
+            )
+        output_path = getattr(task, "result", None)
+        if not output_path:
+            try:
+                kwargs = json.loads(task.task_runner_task.task_kwargs.replace("'", '"'))
+                output_path = kwargs.get("output_path")
+            except Exception:
+                pass
+        if not output_path:
+            self.logger.warning(
+                f"User {request.user.id} tried to download missing report {report_id} (no path)"
+            )
+            return Response(
+                {"detail": "Report file not found"}, status=status.HTTP_404_NOT_FOUND
+            )
+        if output_path.startswith("s3://"):
+            import re
+
+            match = re.match(r"s3://([^/]+)/(.+)", output_path)
+            if not match:
+                return Response(
+                    {"detail": "Invalid S3 URI"},
+                    status=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                )
+            bucket, key = match.groups()
+            s3 = boto3.client("s3")
+            try:
+                with tempfile.NamedTemporaryFile(delete=False) as tmp:
+                    s3.download_fileobj(bucket, key, tmp)
+                    tmp_path = tmp.name
+                with open(tmp_path, "rb") as f:
+                    pdf_data = f.read()
+                os.unlink(tmp_path)
+            except Exception as e:
+                self.logger.warning(
+                    f"User {request.user.id} failed S3 download for report {report_id}: {e}"
+                )
+                return Response(
+                    {"detail": "Report file not found in S3"},
+                    status=status.HTTP_404_NOT_FOUND,
+                )
+            filename = os.path.basename(key)
+        else:
+            if not os.path.exists(output_path):
+                self.logger.warning(
+                    f"User {request.user.id} tried to download missing report {report_id} at {output_path}"
+                )
+                return Response(
+                    {"detail": "Report file not found"},
+                    status=status.HTTP_404_NOT_FOUND,
+                )
+            with open(output_path, "rb") as f:
+                pdf_data = f.read()
+            filename = os.path.basename(output_path)
+        self.logger.info(
+            f"User {request.user.id} downloaded ThreatScore report {report_id} from {output_path}"
+        )
+        response = Response(pdf_data, content_type="application/pdf")
+        response["Content-Disposition"] = f'attachment; filename="{filename}"'
+        return response

api/src/backend/tasks/tasks.py

@@ -1,8 +1,9 @@
+import os
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
 from shutil import rmtree
 
-from celery import chain, shared_task
+from celery import chain, current_task, 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
@@ -33,6 +34,9 @@ from api.v1.serializers import ScanTaskSerializer
 from prowler.lib.check.compliance_models import Compliance
 from prowler.lib.outputs.compliance.generic.generic import GenericCompliance
 from prowler.lib.outputs.finding import Finding as FindingOutput
+from util.compliance_report.threatscore_report_generator import (
+    generate_threatscore_report,
+)
 
 logger = get_task_logger(__name__)
 
@@ -420,3 +424,49 @@ def check_lighthouse_connection_task(lighthouse_config_id: str, tenant_id: str =
             - 'available_models' (list): List of available models if connection is successful.
     """
     return check_lighthouse_connection(lighthouse_config_id=lighthouse_config_id)
+
+
+@shared_task(
+    base=RLSTask,
+    name="scan-threatscore-report",
+    queue="scan-reports",
+)
+@set_tenant(keep_tenant=True)
+def generate_threatscore_report_task(
+    scan_id: str,
+    compliance_id: str,
+    output_path: str,
+    provider_id: str,
+    tenant_id: str,
+    only_failed: bool = True,
+    min_risk_level: int = 4,
+):
+    generate_threatscore_report(
+        scan_id=scan_id,
+        compliance_id=compliance_id,
+        output_path=output_path,
+        provider_id=provider_id,
+        tenant_id=tenant_id,
+        only_failed=only_failed,
+        min_risk_level=min_risk_level,
+    )
+
+    s3_uri = None
+    if os.path.exists(output_path):
+        try:
+            s3_uri = _upload_to_s3(tenant_id, output_path, scan_id)
+        except Exception as e:
+            logger.error(f"Error uploading PDF to S3: {e}")
+
+    from api.models import Task
+
+    task_id = current_task.request.id if hasattr(current_task, "request") else None
+    if task_id:
+        try:
+            task = Task.objects.get(id=task_id)
+            result_path = s3_uri if s3_uri else output_path
+            task.result = result_path
+            task.save(update_fields=["result"])
+        except Exception as e:
+            logger.error(f"Error saving PDF location to Task: {e}")
+    return {"pdf_location": s3_uri if s3_uri else output_path}

util/compliance_report/threatscore_report_generator.py

@@ -2,7 +2,6 @@ import io
 import os
 
 import matplotlib.pyplot as plt
-import requests
 from reportlab.lib import colors
 from reportlab.lib.enums import TA_CENTER
 from reportlab.lib.pagesizes import letter
@@ -21,6 +20,11 @@ from reportlab.platypus import (
     TableStyle,
 )
 
+from api.models import Finding, Provider
+from api.utils import initialize_prowler_provider
+from prowler.lib.check.compliance_models import Compliance
+from prowler.lib.outputs.finding import Finding as FindingOutput
+
 pdfmetrics.registerFont(
     TTFont(
         "PlusJakartaSans",
@@ -42,24 +46,20 @@ def generate_threatscore_report(
     scan_id: str,
     compliance_id: str,
     output_path: str,
-    email: str,
-    password: str,
-    token: str,
-    base_url: str,
+    provider_id: str,
+    tenant_id: str,
     only_failed: bool = True,
     min_risk_level: int = 4,
 ):
     """
-    Generate a PDF compliance report based on Prowler endpoints.
+    Generate a PDF compliance report based on Prowler ORM objects.
 
     Parameters:
     - scan_id: ID of the scan executed by Prowler.
     - compliance_id: ID of the compliance framework (e.g., "nis2_azure").
     - output_path: Output PDF file path (e.g., "threatscore_report.pdf").
-    - email: Email for the API authentication.
-    - password: Password for the API.
-    - token: Token for the API.
-    - base_url: Base URL for the API.
+    - provider_id: Provider ID for the scan.
+    - tenant_id: Tenant ID for the scan.
     - only_failed: If True, only requirements with status "FAIL" will be included in the list of requirements.
     - min_risk_level: Minimum risk level for critical failed requirements.
     """
@@ -136,52 +136,61 @@ def generate_threatscore_report(
         textColor=colors.Color(0.2, 0.2, 0.2),
         fontName="PlusJakartaSans",
     )
-    if not token:
-        if not email or not password:
-            raise Exception("Email and password are required to generate a token")
-        url_credentials = f"{base_url}/api/v1/tokens"
-        payload = {
-            "data": {
-                "type": "tokens",
+
+    provider_obj = Provider.objects.get(id=provider_id)
+    prowler_provider = initialize_prowler_provider(provider_obj)
+    provider_type = provider_obj.provider
+
+    frameworks_bulk = Compliance.get_bulk(provider_type)
+    compliance_obj = frameworks_bulk[compliance_id]
+    compliance_name = compliance_obj.name
+    compliance_version = getattr(compliance_obj, "Version", "N/A")
+    compliance_description = getattr(compliance_obj, "Description", "")
+
+    findings_qs = Finding.all_objects.filter(scan_id=scan_id).order_by("uid")
+    findings = [
+        FindingOutput.transform_api_finding(f, prowler_provider) for f in findings_qs
+    ]
+
+    attrs_map = {}
+    resp_reqs = []
+    resp_attrs = []
+    for req_id, req in compliance_obj.requirements.items():
+        attrs_map[req_id] = {
+            "id": req_id,
+            "attributes": {
+                "metadata": [req.metadata],
+                "check_ids": req.check_ids,
+            },
+            "description": getattr(req, "description", ""),
+        }
+        resp_attrs.append(
+            {
+                "id": req_id,
                 "attributes": {
-                    "email": email,
-                    "password": password,
+                    "framework_description": compliance_description,
+                    "attributes": {"metadata": [req.metadata]},
                 },
             }
-        }
-        resp_credentials = requests.post(
-            url_credentials,
-            json=payload,
-            headers={"Content-Type": "application/vnd.api+json"},
-        ).json()
-        if resp_credentials.get("errors"):
-            print(resp_credentials.get("errors"))
-            raise Exception(resp_credentials.get("errors"))
-        token = resp_credentials.get("data", {}).get("attributes", {}).get("access")
-
-    try:
-        url_reqs = f"{base_url}/api/v1/compliance-overviews/requirements?filter[compliance_id]={compliance_id}&filter[scan_id]={scan_id}"
-        resp_reqs = (
-            requests.get(url_reqs, headers={"Authorization": f"Bearer {token}"})
-            .json()
-            .get("data", [])
         )
-
-        url_attrs = f"{base_url}/api/v1/compliance-overviews/attributes?filter[compliance_id]={compliance_id}"
-        resp_attrs = (
-            requests.get(url_attrs, headers={"Authorization": f"Bearer {token}"})
-            .json()
-            .get("data", [])
+        status = "UNKNOWN"
+        description = getattr(req, "description", "")
+        for f in findings:
+            if f["check_id"] in req.check_ids:
+                status = f["status"]
+                description = f.get("description", description)
+                break
+        resp_reqs.append(
+            {
+                "id": req_id,
+                "attributes": {
+                    "framework": compliance_name,
+                    "version": compliance_version,
+                    "status": status,
+                    "description": description,
+                },
+            }
         )
-    except Exception as e:
-        print(e)
-        raise Exception(e)
-
-    compliance_name = resp_reqs[0]["attributes"]["framework"]
-    compliance_version = resp_reqs[0]["attributes"]["version"]
-    compliance_description = resp_attrs[0]["attributes"]["framework_description"]
-
-    attrs_map = {item["id"]: item["attributes"] for item in resp_attrs}
 
     def create_risk_component(risk_level, weight, score=0):
         """Create a visual risk component similar to the UI design"""
@@ -385,13 +394,7 @@ def generate_threatscore_report(
         return buffer
 
     def get_finding_info(check_id: str):
-        url_find = f"{base_url}/api/v1/findings?filter[check_id]={check_id}&filter[scan]={scan_id}"
-        value = (
-            requests.get(url_find, headers={"Authorization": f"Bearer {token}"})
-            .json()
-            .get("data", [])
-        )
-        return value
+        return [f for f in findings if f["check_id"] == check_id]
 
     doc = SimpleDocTemplate(
         output_path,
@@ -877,23 +880,13 @@ if __name__ == "__main__":
     parser.add_argument(
         "--output", default="threatscore_report.pdf", help="Output PDF file path"
     )
-    parser.add_argument("--email", required=True, help="Email for the API")
-    parser.add_argument("--password", required=True, help="Password for the API")
+    parser.add_argument("--provider-id", required=True, help="Provider ID")
+    parser.add_argument("--tenant-id", required=True, help="Tenant ID")
     parser.add_argument(
         "--only-failed",
         action="store_true",
         help="Only include failed requirements in the list of requirements",
     )
-    parser.add_argument(
-        "--token",
-        default="",
-        help="Token for the API",
-    )
-    parser.add_argument(
-        "--base-url",
-        default="http://localhost:8080",
-        help="Base URL for the API",
-    )
     parser.add_argument(
         "--min-risk-level",
         type=int,
@@ -906,10 +899,8 @@ if __name__ == "__main__":
         args.scan_id,
         args.compliance_id,
         args.output,
-        args.email,
-        args.password,
-        args.token,
-        args.base_url,
+        args.provider_id,
+        args.tenant_id,
         args.only_failed,
         args.min_risk_level,
     )