diff --git a/api/src/backend/api/v1/views.py b/api/src/backend/api/v1/views.py index 4f7075ff78..dfdda6391d 100644 --- a/api/src/backend/api/v1/views.py +++ b/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[^/]+)") + 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[^/]+)") + 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 diff --git a/api/src/backend/tasks/tasks.py b/api/src/backend/tasks/tasks.py index addb2c247e..fe0ba11a46 100644 --- a/api/src/backend/tasks/tasks.py +++ b/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} diff --git a/util/compliance_report/threatscore_report_generator.py b/util/compliance_report/threatscore_report_generator.py index 0ea49aa868..d740d06ec1 100644 --- a/util/compliance_report/threatscore_report_generator.py +++ b/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, )