diff --git a/.github/workflows/api-pull-request.yml b/.github/workflows/api-pull-request.yml index 675c0ae500..604c878bf6 100644 --- a/.github/workflows/api-pull-request.yml +++ b/.github/workflows/api-pull-request.yml @@ -85,6 +85,14 @@ jobs: api/README.md api/mkdocs.yml + - name: Replace @master with current branch in pyproject.toml + working-directory: ./api + if: steps.are-non-ignored-files-changed.outputs.any_changed == 'true' + run: | + BRANCH_NAME="${GITHUB_HEAD_REF:-${GITHUB_REF_NAME}}" + echo "Using branch: $BRANCH_NAME" + sed -i "s|@master|@$BRANCH_NAME|g" pyproject.toml + - name: Install poetry working-directory: ./api if: steps.are-non-ignored-files-changed.outputs.any_changed == 'true' @@ -92,6 +100,12 @@ jobs: python -m pip install --upgrade pip pipx install poetry==2.1.1 + - name: Update poetry.lock after the branch name change + working-directory: ./api + if: steps.are-non-ignored-files-changed.outputs.any_changed == 'true' + run: | + poetry lock + - name: Set up Python ${{ matrix.python-version }} if: steps.are-non-ignored-files-changed.outputs.any_changed == 'true' uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 diff --git a/api/CHANGELOG.md b/api/CHANGELOG.md index a8671b5478..2b20a6ea5b 100644 --- a/api/CHANGELOG.md +++ b/api/CHANGELOG.md @@ -8,6 +8,8 @@ All notable changes to the **Prowler API** are documented in this file. ### Added - Added M365 as a new provider [(#7563)](https://github.com/prowler-cloud/prowler/pull/7563). +- Added a `compliance/` folder and ZIP‐export functionality for all compliance reports.[(#7653)](https://github.com/prowler-cloud/prowler/pull/7653). +- Added a new API endpoint to fetch and download any specific compliance file by name [(#7653)](https://github.com/prowler-cloud/prowler/pull/7653). --- diff --git a/api/poetry.lock b/api/poetry.lock index f9a6ec72ed..94cd64b994 100644 --- a/api/poetry.lock +++ b/api/poetry.lock @@ -1469,14 +1469,14 @@ with-social = ["django-allauth[socialaccount] (>=64.0.0)"] [[package]] name = "django" -version = "5.1.7" +version = "5.1.8" description = "A high-level Python web framework that encourages rapid development and clean, pragmatic design." optional = false python-versions = ">=3.10" groups = ["main", "dev"] files = [ - {file = "Django-5.1.7-py3-none-any.whl", hash = "sha256:1323617cb624add820cb9611cdcc788312d250824f92ca6048fda8625514af2b"}, - {file = "Django-5.1.7.tar.gz", hash = "sha256:30de4ee43a98e5d3da36a9002f287ff400b43ca51791920bfb35f6917bfe041c"}, + {file = "Django-5.1.8-py3-none-any.whl", hash = "sha256:11b28fa4b00e59d0def004e9ee012fefbb1065a5beb39ee838983fd24493ad4f"}, + {file = "Django-5.1.8.tar.gz", hash = "sha256:42e92a1dd2810072bcc40a39a212b693f94406d0ba0749e68eb642f31dc770b4"}, ] [package.dependencies] @@ -5488,4 +5488,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.1" python-versions = ">=3.11,<3.13" -content-hash = "6e196101601537d1bd6f2eb400eccbc09e247d8e2a625843733e735e03ea6d9b" +content-hash = "fd475f6b8f0b5c64e9f4e1d3e76ddc1d53a2caa5b028ee2b4ed9c8b74d3c5627" diff --git a/api/pyproject.toml b/api/pyproject.toml index f642e5b87d..b3da78cc70 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -7,7 +7,7 @@ authors = [{name = "Prowler Engineering", email = "engineering@prowler.com"}] dependencies = [ "celery[pytest] (>=5.4.0,<6.0.0)", "dj-rest-auth[with_social,jwt] (==7.0.1)", - "django==5.1.7", + "django==5.1.8", "django-allauth==65.4.1", "django-celery-beat (>=2.7.0,<3.0.0)", "django-celery-results (>=2.5.1,<3.0.0)", diff --git a/api/src/backend/api/compliance.py b/api/src/backend/api/compliance.py index ae1b89a4a9..96b5e313f3 100644 --- a/api/src/backend/api/compliance.py +++ b/api/src/backend/api/compliance.py @@ -1,12 +1,38 @@ from types import MappingProxyType +from api.models import Provider +from prowler.config.config import get_available_compliance_frameworks from prowler.lib.check.compliance_models import Compliance from prowler.lib.check.models import CheckMetadata -from api.models import Provider - PROWLER_COMPLIANCE_OVERVIEW_TEMPLATE = {} PROWLER_CHECKS = {} +AVAILABLE_COMPLIANCE_FRAMEWORKS = {} + + +def get_compliance_frameworks(provider_type: Provider.ProviderChoices) -> list[str]: + """ + Retrieve and cache the list of available compliance frameworks for a specific cloud provider. + + This function lazily loads and caches the available compliance frameworks (e.g., CIS, MITRE, ISO) + for each provider type (AWS, Azure, GCP, etc.) on first access. Subsequent calls for the same + provider will return the cached result. + + Args: + provider_type (Provider.ProviderChoices): The cloud provider type for which to retrieve + available compliance frameworks (e.g., "aws", "azure", "gcp", "m365"). + + Returns: + list[str]: A list of framework identifiers (e.g., "cis_1.4_aws", "mitre_attack_azure") available + for the given provider. + """ + global AVAILABLE_COMPLIANCE_FRAMEWORKS + if provider_type not in AVAILABLE_COMPLIANCE_FRAMEWORKS: + AVAILABLE_COMPLIANCE_FRAMEWORKS[provider_type] = ( + get_available_compliance_frameworks(provider_type) + ) + + return AVAILABLE_COMPLIANCE_FRAMEWORKS[provider_type] def get_prowler_provider_checks(provider_type: Provider.ProviderChoices): diff --git a/api/src/backend/api/specs/v1.yaml b/api/src/backend/api/specs/v1.yaml index bb7c580815..962ec773c2 100644 --- a/api/src/backend/api/specs/v1.yaml +++ b/api/src/backend/api/specs/v1.yaml @@ -4503,6 +4503,47 @@ paths: schema: $ref: '#/components/schemas/ScanUpdateResponse' description: '' + /api/v1/scans/{id}/compliance/{name}: + get: + operationId: scan_compliance_download + description: Download a specific compliance report (e.g., 'cis_1.4_aws') as + a CSV file. + summary: Retrieve compliance report as CSV + parameters: + - in: query + name: fields[scan-reports] + schema: + type: array + items: + type: string + enum: + - id + - name + description: endpoint return only specific fields in the response on a per-type + basis by including a fields[TYPE] query parameter. + explode: false + - in: path + name: id + schema: + type: string + format: uuid + description: A UUID string identifying this scan. + required: true + - in: path + name: name + schema: + type: string + description: The compliance report name, like 'cis_1.4_aws' + required: true + tags: + - Scan + security: + - jwtAuth: [] + responses: + '200': + description: CSV file containing the compliance report + '404': + description: Compliance report not found /api/v1/scans/{id}/report: get: operationId: scans_report_retrieve diff --git a/api/src/backend/api/tests/test_views.py b/api/src/backend/api/tests/test_views.py index 002af353da..455f9f98af 100644 --- a/api/src/backend/api/tests/test_views.py +++ b/api/src/backend/api/tests/test_views.py @@ -3,16 +3,17 @@ import io import json import os from datetime import datetime, timedelta, timezone -from unittest.mock import ANY, Mock, patch +from unittest.mock import ANY, MagicMock, Mock, patch import jwt import pytest -from botocore.exceptions import NoCredentialsError +from botocore.exceptions import ClientError, NoCredentialsError from conftest import API_JSON_CONTENT_TYPE, TEST_PASSWORD, TEST_USER from django.conf import settings from django.urls import reverse from rest_framework import status +from api.compliance import get_compliance_frameworks from api.models import ( ComplianceOverview, Integration, @@ -2277,7 +2278,8 @@ class TestScanViewSet: scan.save() monkeypatch.setattr( - "api.v1.views.env", type("env", (), {"str": lambda self, key: bucket})() + "api.v1.views.env", + type("env", (), {"str": lambda self, *args, **kwargs: "test-bucket"})(), ) class FakeS3Client: @@ -2346,6 +2348,263 @@ class TestScanViewSet: assert content_disposition.startswith('attachment; filename="') assert f'filename="{file_path.name}"' in content_disposition + def test_compliance_invalid_framework(self, authenticated_client, scans_fixture): + scan = scans_fixture[0] + scan.state = StateChoices.COMPLETED + scan.output_location = "dummy" + scan.save() + + url = reverse("scan-compliance", kwargs={"pk": scan.id, "name": "invalid"}) + resp = authenticated_client.get(url) + assert resp.status_code == status.HTTP_404_NOT_FOUND + assert resp.json()["errors"]["detail"] == "Compliance 'invalid' not found." + + def test_compliance_executing( + self, authenticated_client, scans_fixture, monkeypatch + ): + scan = scans_fixture[0] + scan.state = StateChoices.EXECUTING + scan.save() + task = Task.objects.create(tenant_id=scan.tenant_id) + scan.task = task + scan.save() + dummy = {"id": str(task.id), "state": StateChoices.EXECUTING} + + monkeypatch.setattr( + "api.v1.views.TaskSerializer", + lambda *args, **kwargs: type("S", (), {"data": dummy}), + ) + + framework = get_compliance_frameworks(scan.provider.provider)[0] + url = reverse("scan-compliance", kwargs={"pk": scan.id, "name": framework}) + resp = authenticated_client.get(url) + assert resp.status_code == status.HTTP_202_ACCEPTED + assert "Content-Location" in resp + assert dummy["id"] in resp["Content-Location"] + + def test_compliance_no_output(self, authenticated_client, scans_fixture): + scan = scans_fixture[0] + scan.state = StateChoices.COMPLETED + scan.output_location = "" + scan.save() + + framework = get_compliance_frameworks(scan.provider.provider)[0] + url = reverse("scan-compliance", kwargs={"pk": scan.id, "name": framework}) + resp = authenticated_client.get(url) + assert resp.status_code == status.HTTP_404_NOT_FOUND + assert resp.json()["errors"]["detail"] == "The scan has no reports." + + def test_compliance_s3_no_credentials( + self, authenticated_client, scans_fixture, monkeypatch + ): + scan = scans_fixture[0] + bucket = "bucket" + key = "file.zip" + scan.output_location = f"s3://{bucket}/{key}" + scan.state = StateChoices.COMPLETED + scan.save() + + monkeypatch.setattr( + "api.v1.views.get_s3_client", + lambda: (_ for _ in ()).throw(NoCredentialsError()), + ) + + framework = get_compliance_frameworks(scan.provider.provider)[0] + url = reverse("scan-compliance", kwargs={"pk": scan.id, "name": framework}) + resp = authenticated_client.get(url) + assert resp.status_code == status.HTTP_403_FORBIDDEN + assert resp.json()["errors"]["detail"] == "There is a problem with credentials." + + def test_compliance_s3_success( + self, authenticated_client, scans_fixture, monkeypatch + ): + scan = scans_fixture[0] + bucket = "bucket" + prefix = "path/scan.zip" + scan.output_location = f"s3://{bucket}/{prefix}" + scan.state = StateChoices.COMPLETED + scan.save() + + monkeypatch.setattr( + "api.v1.views.env", + type("env", (), {"str": lambda self, *args, **kwargs: "test-bucket"})(), + ) + + match_key = "path/compliance/mitre_attack_aws.csv" + + class FakeS3Client: + def list_objects_v2(self, Bucket, Prefix): + return {"Contents": [{"Key": match_key}]} + + def get_object(self, Bucket, Key): + return {"Body": io.BytesIO(b"ignored")} + + monkeypatch.setattr("api.v1.views.get_s3_client", lambda: FakeS3Client()) + + framework = match_key.split("/")[-1].split(".")[0] + url = reverse("scan-compliance", kwargs={"pk": scan.id, "name": framework}) + resp = authenticated_client.get(url) + assert resp.status_code == status.HTTP_200_OK + cd = resp["Content-Disposition"] + assert cd.startswith('attachment; filename="') + assert cd.endswith('filename="mitre_attack_aws.csv"') + + def test_compliance_s3_not_found( + self, authenticated_client, scans_fixture, monkeypatch + ): + scan = scans_fixture[0] + bucket = "bucket" + scan.output_location = f"s3://{bucket}/x/scan.zip" + scan.state = StateChoices.COMPLETED + scan.save() + + monkeypatch.setattr( + "api.v1.views.env", + type("env", (), {"str": lambda self, *args, **kwargs: "test-bucket"})(), + ) + + class FakeS3Client: + def list_objects_v2(self, Bucket, Prefix): + return {"Contents": []} + + def get_object(self, Bucket, Key): + return {"Body": io.BytesIO(b"ignored")} + + monkeypatch.setattr("api.v1.views.get_s3_client", lambda: FakeS3Client()) + + url = reverse("scan-compliance", kwargs={"pk": scan.id, "name": "cis_1.4_aws"}) + resp = authenticated_client.get(url) + assert resp.status_code == status.HTTP_404_NOT_FOUND + assert ( + resp.json()["errors"]["detail"] + == "No compliance file found for name 'cis_1.4_aws'." + ) + + def test_compliance_local_file( + self, authenticated_client, scans_fixture, tmp_path, monkeypatch + ): + scan = scans_fixture[0] + scan.state = StateChoices.COMPLETED + base = tmp_path / "reports" + comp_dir = base / "compliance" + comp_dir.mkdir(parents=True) + fname = comp_dir / "scan_cis.csv" + fname.write_bytes(b"ignored") + + scan.output_location = str(base / "scan.zip") + scan.save() + + monkeypatch.setattr( + glob, + "glob", + lambda p: [str(fname)] if p.endswith("*_cis_1.4_aws.csv") else [], + ) + + url = reverse("scan-compliance", kwargs={"pk": scan.id, "name": "cis_1.4_aws"}) + resp = authenticated_client.get(url) + assert resp.status_code == status.HTTP_200_OK + cd = resp["Content-Disposition"] + assert cd.startswith('attachment; filename="') + assert cd.endswith(f'filename="{fname.name}"') + + @patch("api.v1.views.Task.objects.get") + @patch("api.v1.views.TaskSerializer") + def test__get_task_status_returns_none_if_task_not_executing( + self, mock_task_serializer, mock_task_get, authenticated_client, scans_fixture + ): + scan = scans_fixture[0] + scan.state = StateChoices.COMPLETED + scan.output_location = "dummy" + scan.save() + + task = Task.objects.create(tenant_id=scan.tenant_id) + mock_task_get.return_value = task + mock_task_serializer.return_value.data = { + "id": str(task.id), + "state": StateChoices.COMPLETED, + } + + url = reverse("scan-report", kwargs={"pk": scan.id}) + response = authenticated_client.get(url) + + assert response.status_code == status.HTTP_404_NOT_FOUND + + @patch("api.v1.views.get_s3_client") + @patch("api.v1.views.sentry_sdk.capture_exception") + def test_compliance_list_objects_client_error( + self, + mock_sentry_capture, + mock_get_s3_client, + authenticated_client, + scans_fixture, + ): + scan = scans_fixture[0] + scan.output_location = "s3://test-bucket/path/to/scan.zip" + scan.state = StateChoices.COMPLETED + scan.save() + + fake_client = MagicMock() + fake_client.list_objects_v2.side_effect = ClientError( + {"Error": {"Code": "InternalError"}}, "ListObjectsV2" + ) + mock_get_s3_client.return_value = fake_client + + framework = get_compliance_frameworks(scan.provider.provider)[0] + url = reverse("scan-compliance", kwargs={"pk": scan.id, "name": framework}) + response = authenticated_client.get(url) + + assert response.status_code == status.HTTP_502_BAD_GATEWAY + assert ( + response.json()["errors"]["detail"] + == "Unable to list compliance files in S3: encountered an AWS error." + ) + mock_sentry_capture.assert_called() + + @patch("api.v1.views.get_s3_client") + def test_report_s3_nosuchkey( + self, mock_get_s3_client, authenticated_client, scans_fixture + ): + scan = scans_fixture[0] + scan.output_location = "s3://test-bucket/report.zip" + scan.state = StateChoices.COMPLETED + scan.save() + + fake_client = MagicMock() + fake_client.get_object.side_effect = ClientError( + {"Error": {"Code": "NoSuchKey"}}, "GetObject" + ) + mock_get_s3_client.return_value = fake_client + + url = reverse("scan-report", kwargs={"pk": scan.id}) + response = authenticated_client.get(url) + + assert response.status_code == status.HTTP_404_NOT_FOUND + assert response.json()["errors"]["detail"] == "The scan has no reports." + + @patch("api.v1.views.get_s3_client") + def test_report_s3_client_error_other( + self, mock_get_s3_client, authenticated_client, scans_fixture + ): + scan = scans_fixture[0] + scan.output_location = "s3://test-bucket/report.zip" + scan.state = StateChoices.COMPLETED + scan.save() + + fake_client = MagicMock() + fake_client.get_object.side_effect = ClientError( + {"Error": {"Code": "AccessDenied"}}, "GetObject" + ) + mock_get_s3_client.return_value = fake_client + + url = reverse("scan-report", kwargs={"pk": scan.id}) + response = authenticated_client.get(url) + + assert response.status_code == status.HTTP_403_FORBIDDEN + assert ( + response.json()["errors"]["detail"] + == "There is a problem with credentials." + ) + @pytest.mark.django_db class TestTaskViewSet: diff --git a/api/src/backend/api/v1/serializers.py b/api/src/backend/api/v1/serializers.py index 920384ba65..c0440f3ef0 100644 --- a/api/src/backend/api/v1/serializers.py +++ b/api/src/backend/api/v1/serializers.py @@ -960,6 +960,15 @@ class ScanReportSerializer(serializers.Serializer): fields = ["id"] +class ScanComplianceReportSerializer(serializers.Serializer): + id = serializers.CharField(source="scan") + name = serializers.CharField() + + class Meta: + resource_name = "scan-reports" + fields = ["id", "name"] + + class ResourceTagSerializer(RLSSerializer): """ Serializer for the ResourceTag model diff --git a/api/src/backend/api/v1/views.py b/api/src/backend/api/v1/views.py index 606078b558..c0a2e824fe 100644 --- a/api/src/backend/api/v1/views.py +++ b/api/src/backend/api/v1/views.py @@ -55,6 +55,7 @@ from tasks.tasks import ( ) from api.base_views import BaseRLSViewSet, BaseTenantViewset, BaseUserViewset +from api.compliance import get_compliance_frameworks from api.db_router import MainRouter from api.filters import ( ComplianceOverviewFilter, @@ -134,6 +135,7 @@ from api.v1.serializers import ( RoleProviderGroupRelationshipSerializer, RoleSerializer, RoleUpdateSerializer, + ScanComplianceReportSerializer, ScanCreateSerializer, ScanReportSerializer, ScanSerializer, @@ -1150,6 +1152,27 @@ class ProviderViewSet(BaseRLSViewSet): 404: OpenApiResponse(description="The scan has no reports"), }, ), + compliance=extend_schema( + tags=["Scan"], + summary="Retrieve compliance report as CSV", + description="Download a specific compliance report (e.g., 'cis_1.4_aws') as a CSV file.", + parameters=[ + OpenApiParameter( + name="name", + type=str, + location=OpenApiParameter.PATH, + required=True, + description="The compliance report name, like 'cis_1.4_aws'", + ), + ], + responses={ + 200: OpenApiResponse( + description="CSV file containing the compliance report" + ), + 404: OpenApiResponse(description="Compliance report not found"), + }, + request=None, + ), ) @method_decorator(CACHE_DECORATOR, name="list") @method_decorator(CACHE_DECORATOR, name="retrieve") @@ -1202,6 +1225,10 @@ class ScanViewSet(BaseRLSViewSet): if hasattr(self, "response_serializer_class"): return self.response_serializer_class return ScanReportSerializer + elif self.action == "compliance": + if hasattr(self, "response_serializer_class"): + return self.response_serializer_class + return ScanComplianceReportSerializer return super().get_serializer_class() def partial_update(self, request, *args, **kwargs): @@ -1219,70 +1246,111 @@ class ScanViewSet(BaseRLSViewSet): ) return Response(data=read_serializer.data, status=status.HTTP_200_OK) - @action(detail=True, methods=["get"], url_name="report") - def report(self, request, pk=None): - scan_instance = self.get_object() + def _get_task_status(self, scan_instance): + """ + Returns task status if the scan or its associated report-generation task is still executing. - if scan_instance.state == StateChoices.EXECUTING: - # If the scan is still running, return the task - prowler_task = Task.objects.get(id=scan_instance.task.id) - self.response_serializer_class = TaskSerializer - output_serializer = self.get_serializer(prowler_task) - return Response( - data=output_serializer.data, - status=status.HTTP_202_ACCEPTED, - headers={ - "Content-Location": reverse( - "task-detail", kwargs={"pk": output_serializer.data["id"]} - ) - }, - ) + If the scan is in an EXECUTING state or if a background task related to report generation + is found and also executing, this method returns a 202 Accepted response with the task + metadata and a `Content-Location` header pointing to the task detail endpoint. - try: - output_celery_task = Task.objects.get( - task_runner_task__task_name="scan-report", - task_runner_task__task_args__contains=pk, - ) - self.response_serializer_class = TaskSerializer - output_serializer = self.get_serializer(output_celery_task) - if output_serializer.data["state"] == StateChoices.EXECUTING: - # If the task is still running, return the task - return Response( - data=output_serializer.data, - status=status.HTTP_202_ACCEPTED, - headers={ - "Content-Location": reverse( - "task-detail", kwargs={"pk": output_serializer.data["id"]} - ) - }, - ) - except Task.DoesNotExist: - # If the task does not exist, it means that the task is removed from the database - pass + Args: + scan_instance (Scan): The scan instance for which the task status is being checked. - output_location = scan_instance.output_location - if not output_location: - return Response( - {"detail": "The scan has no reports."}, - status=status.HTTP_404_NOT_FOUND, - ) + Returns: + Response or None: + - A `Response` with HTTP 202 status and serialized task data if the task is executing. + - `None` if no running task is found or if the task has already completed. + """ + task = None - if scan_instance.output_location.startswith("s3://"): + if scan_instance.state == StateChoices.EXECUTING and scan_instance.task: + task = scan_instance.task + else: try: - s3_client = get_s3_client() + task = Task.objects.get( + task_runner_task__task_name="scan-report", + task_runner_task__task_args__contains=str(scan_instance.id), + ) + except Task.DoesNotExist: + return None + + self.response_serializer_class = TaskSerializer + serializer = self.get_serializer(task) + + if serializer.data.get("state") != StateChoices.EXECUTING: + return None + + return Response( + data=serializer.data, + status=status.HTTP_202_ACCEPTED, + headers={ + "Content-Location": reverse( + "task-detail", kwargs={"pk": serializer.data["id"]} + ) + }, + ) + + def _load_file(self, path_pattern, s3=False, bucket=None, list_objects=False): + """ + Loads a binary file (e.g., ZIP or CSV) and returns its content and filename. + + Depending on the input parameters, this method supports loading: + - From S3 using a direct key. + - From S3 by listing objects under a prefix and matching suffix. + - From the local filesystem using glob pattern matching. + + Args: + path_pattern (str): The key or glob pattern representing the file location. + s3 (bool, optional): Whether the file is stored in S3. Defaults to False. + bucket (str, optional): The name of the S3 bucket, required if `s3=True`. Defaults to None. + list_objects (bool, optional): If True and `s3=True`, list objects by prefix to find the file. Defaults to False. + + Returns: + tuple[bytes, str]: A tuple containing the file content as bytes and the filename if successful. + Response: A DRF `Response` object with an appropriate status and error detail if an error occurs. + """ + if s3: + try: + client = get_s3_client() except (ClientError, NoCredentialsError, ParamValidationError): return Response( {"detail": "There is a problem with credentials."}, status=status.HTTP_403_FORBIDDEN, ) - - bucket_name = env.str("DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET") - key = output_location[len(f"s3://{bucket_name}/") :] + if list_objects: + # list keys under prefix then match suffix + prefix = os.path.dirname(path_pattern) + suffix = os.path.basename(path_pattern) + try: + resp = client.list_objects_v2(Bucket=bucket, Prefix=prefix) + except ClientError as e: + sentry_sdk.capture_exception(e) + return Response( + { + "detail": "Unable to list compliance files in S3: encountered an AWS error." + }, + status=status.HTTP_502_BAD_GATEWAY, + ) + contents = resp.get("Contents", []) + keys = [obj["Key"] for obj in contents if obj["Key"].endswith(suffix)] + if not keys: + return Response( + { + "detail": f"No compliance file found for name '{os.path.splitext(suffix)[0]}'." + }, + status=status.HTTP_404_NOT_FOUND, + ) + # path_pattern here is prefix, but in compliance we build correct suffix check before + key = keys[0] + else: + # path_pattern is exact key + key = path_pattern try: - s3_object = s3_client.get_object(Bucket=bucket_name, Key=key) + s3_obj = client.get_object(Bucket=bucket, Key=key) except ClientError as e: - error_code = e.response.get("Error", {}).get("Code") - if error_code == "NoSuchKey": + code = e.response.get("Error", {}).get("Code") + if code == "NoSuchKey": return Response( {"detail": "The scan has no reports."}, status=status.HTTP_404_NOT_FOUND, @@ -1291,28 +1359,97 @@ class ScanViewSet(BaseRLSViewSet): {"detail": "There is a problem with credentials."}, status=status.HTTP_403_FORBIDDEN, ) - file_content = s3_object["Body"].read() - filename = os.path.basename(output_location.split("/")[-1]) + content = s3_obj["Body"].read() + filename = os.path.basename(key) else: - zip_files = glob.glob(output_location) - try: - file_path = zip_files[0] - except IndexError as e: - sentry_sdk.capture_exception(e) + files = glob.glob(path_pattern) + if not files: return Response( {"detail": "The scan has no reports."}, status=status.HTTP_404_NOT_FOUND, ) - with open(file_path, "rb") as f: - file_content = f.read() - filename = os.path.basename(file_path) + filepath = files[0] + with open(filepath, "rb") as f: + content = f.read() + filename = os.path.basename(filepath) - response = HttpResponse( - file_content, content_type="application/x-zip-compressed" - ) + return content, filename + + def _serve_file(self, content, filename, content_type): + response = HttpResponse(content, content_type=content_type) response["Content-Disposition"] = f'attachment; filename="{filename}"' + return response + @action(detail=True, methods=["get"], url_name="report") + def report(self, request, pk=None): + scan = self.get_object() + # Check for executing tasks + 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."}, 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}/") + loader = self._load_file( + key_prefix, s3=True, bucket=bucket, list_objects=False + ) + else: + loader = self._load_file(scan.output_location, s3=False) + + if isinstance(loader, Response): + return loader + + content, filename = loader + return self._serve_file(content, filename, "application/x-zip-compressed") + + @action( + detail=True, + methods=["get"], + url_path="compliance/(?P[^/]+)", + url_name="compliance", + ) + def compliance(self, request, pk=None, name=None): + scan = self.get_object() + if name not in get_compliance_frameworks(scan.provider.provider): + return Response( + {"detail": f"Compliance '{name}' not found."}, + status=status.HTTP_404_NOT_FOUND, + ) + + 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."}, 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"{name}.csv" + ) + 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"*_{name}.csv") + loader = self._load_file(pattern, s3=False) + + if isinstance(loader, Response): + return loader + + content, filename = loader + return self._serve_file(content, filename, "text/csv") + def create(self, request, *args, **kwargs): input_serializer = self.get_serializer(data=request.data) input_serializer.is_valid(raise_exception=True) diff --git a/api/src/backend/tasks/jobs/export.py b/api/src/backend/tasks/jobs/export.py index 11c9e5c4cf..fb405176d5 100644 --- a/api/src/backend/tasks/jobs/export.py +++ b/api/src/backend/tasks/jobs/export.py @@ -13,6 +13,38 @@ from prowler.config.config import ( json_ocsf_file_suffix, output_file_timestamp, ) +from prowler.lib.outputs.compliance.aws_well_architected.aws_well_architected import ( + AWSWellArchitected, +) +from prowler.lib.outputs.compliance.cis.cis_aws import AWSCIS +from prowler.lib.outputs.compliance.cis.cis_azure import AzureCIS +from prowler.lib.outputs.compliance.cis.cis_gcp import GCPCIS +from prowler.lib.outputs.compliance.cis.cis_kubernetes import KubernetesCIS +from prowler.lib.outputs.compliance.cis.cis_m365 import M365CIS +from prowler.lib.outputs.compliance.ens.ens_aws import AWSENS +from prowler.lib.outputs.compliance.ens.ens_azure import AzureENS +from prowler.lib.outputs.compliance.ens.ens_gcp import GCPENS +from prowler.lib.outputs.compliance.iso27001.iso27001_aws import AWSISO27001 +from prowler.lib.outputs.compliance.iso27001.iso27001_azure import AzureISO27001 +from prowler.lib.outputs.compliance.iso27001.iso27001_gcp import GCPISO27001 +from prowler.lib.outputs.compliance.iso27001.iso27001_kubernetes import ( + KubernetesISO27001, +) +from prowler.lib.outputs.compliance.kisa_ismsp.kisa_ismsp_aws import AWSKISAISMSP +from prowler.lib.outputs.compliance.mitre_attack.mitre_attack_aws import AWSMitreAttack +from prowler.lib.outputs.compliance.mitre_attack.mitre_attack_azure import ( + AzureMitreAttack, +) +from prowler.lib.outputs.compliance.mitre_attack.mitre_attack_gcp import GCPMitreAttack +from prowler.lib.outputs.compliance.prowler_threatscore.prowler_threatscore_aws import ( + ProwlerThreatScoreAWS, +) +from prowler.lib.outputs.compliance.prowler_threatscore.prowler_threatscore_azure import ( + ProwlerThreatScoreAzure, +) +from prowler.lib.outputs.compliance.prowler_threatscore.prowler_threatscore_gcp import ( + ProwlerThreatScoreGCP, +) from prowler.lib.outputs.csv.csv import CSV from prowler.lib.outputs.html.html import HTML from prowler.lib.outputs.ocsf.ocsf import OCSF @@ -20,6 +52,43 @@ from prowler.lib.outputs.ocsf.ocsf import OCSF logger = get_task_logger(__name__) +COMPLIANCE_CLASS_MAP = { + "aws": [ + (lambda name: name.startswith("cis_"), AWSCIS), + (lambda name: name == "mitre_attack_aws", AWSMitreAttack), + (lambda name: name.startswith("ens_"), AWSENS), + ( + lambda name: name.startswith("aws_well_architected_framework"), + AWSWellArchitected, + ), + (lambda name: name.startswith("iso27001_"), AWSISO27001), + (lambda name: name.startswith("kisa"), AWSKISAISMSP), + (lambda name: name == "prowler_threatscore_aws", ProwlerThreatScoreAWS), + ], + "azure": [ + (lambda name: name.startswith("cis_"), AzureCIS), + (lambda name: name == "mitre_attack_azure", AzureMitreAttack), + (lambda name: name.startswith("ens_"), AzureENS), + (lambda name: name.startswith("iso27001_"), AzureISO27001), + (lambda name: name == "prowler_threatscore_azure", ProwlerThreatScoreAzure), + ], + "gcp": [ + (lambda name: name.startswith("cis_"), GCPCIS), + (lambda name: name == "mitre_attack_gcp", GCPMitreAttack), + (lambda name: name.startswith("ens_"), GCPENS), + (lambda name: name.startswith("iso27001_"), GCPISO27001), + (lambda name: name == "prowler_threatscore_gcp", ProwlerThreatScoreGCP), + ], + "kubernetes": [ + (lambda name: name.startswith("cis_"), KubernetesCIS), + (lambda name: name.startswith("iso27001_"), KubernetesISO27001), + ], + "m365": [ + (lambda name: name.startswith("cis_"), M365CIS), + ], +} + + # Predefined mapping for output formats and their configurations OUTPUT_FORMATS_MAPPING = { "csv": { @@ -43,13 +112,17 @@ def _compress_output_files(output_directory: str) -> str: str: The full path to the newly created ZIP archive. """ zip_path = f"{output_directory}.zip" + parent_dir = os.path.dirname(output_directory) + zip_path_abs = os.path.abspath(zip_path) 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}", - ) + for foldername, _, filenames in os.walk(parent_dir): + for filename in filenames: + file_path = os.path.join(foldername, filename) + if os.path.abspath(file_path) == zip_path_abs: + continue + arcname = os.path.relpath(file_path, start=parent_dir) + zipf.write(file_path, arcname) return zip_path @@ -102,25 +175,38 @@ def _upload_to_s3(tenant_id: str, zip_path: str, scan_id: str) -> str: Raises: botocore.exceptions.ClientError: If the upload attempt to S3 fails for any reason. """ - if not base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET: - return + bucket = base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET + if not bucket: + return None try: s3 = get_s3_client() - s3_key = f"{tenant_id}/{scan_id}/{os.path.basename(zip_path)}" + + # Upload the ZIP file (outputs) to the S3 bucket + zip_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, + Bucket=bucket, + Key=zip_key, ) - return f"s3://{base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET}/{s3_key}" + + # Upload the compliance directory to the S3 bucket + compliance_dir = os.path.join(os.path.dirname(zip_path), "compliance") + for filename in os.listdir(compliance_dir): + local_path = os.path.join(compliance_dir, filename) + if not os.path.isfile(local_path): + continue + file_key = f"{tenant_id}/{scan_id}/compliance/{filename}" + s3.upload_file(Filename=local_path, Bucket=bucket, Key=file_key) + + return f"s3://{base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET}/{zip_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: +) -> tuple[str, str]: """ Generate a file system path for the output directory of a prowler scan. @@ -145,7 +231,8 @@ def _generate_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' + '/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' """ path = ( f"{output_directory}/{tenant_id}/{scan_id}/prowler-output-" @@ -153,4 +240,10 @@ def _generate_output_directory( ) os.makedirs("/".join(path.split("/")[:-1]), exist_ok=True) - return path + compliance_path = ( + f"{output_directory}/{tenant_id}/{scan_id}/compliance/prowler-output-" + f"{prowler_provider}-{output_file_timestamp}" + ) + os.makedirs("/".join(compliance_path.split("/")[:-1]), exist_ok=True) + + return path, compliance_path diff --git a/api/src/backend/tasks/tasks.py b/api/src/backend/tasks/tasks.py index 5af7895394..fa7012e7d8 100644 --- a/api/src/backend/tasks/tasks.py +++ b/api/src/backend/tasks/tasks.py @@ -10,6 +10,7 @@ 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 ( + COMPLIANCE_CLASS_MAP, OUTPUT_FORMATS_MAPPING, _compress_output_files, _generate_output_directory, @@ -18,11 +19,14 @@ from tasks.jobs.export import ( from tasks.jobs.scan import aggregate_findings, perform_prowler_scan from tasks.utils import batched, get_next_execution_datetime +from api.compliance import get_compliance_frameworks from api.db_utils import rls_transaction from api.decorators import set_tenant from api.models import Finding, Provider, Scan, ScanSummary, StateChoices from api.utils import initialize_prowler_provider 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 logger = get_task_logger(__name__) @@ -251,84 +255,106 @@ def generate_outputs(scan_id: str, provider_id: str, tenant_id: str): logger.info(f"No findings found for scan {scan_id}") return {"upload": False} - # Initialize the prowler provider - prowler_provider = initialize_prowler_provider(Provider.objects.get(id=provider_id)) + provider_obj = Provider.objects.get(id=provider_id) + prowler_provider = initialize_prowler_provider(provider_obj) + provider_uid = provider_obj.uid + provider_type = provider_obj.provider - # 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( + frameworks_bulk = Compliance.get_bulk(provider_type) + frameworks_avail = get_compliance_frameworks(provider_type) + out_dir, comp_dir = _generate_output_directory( DJANGO_TMP_OUTPUT_DIRECTORY, provider_uid, tenant_id, scan_id ) - # Define auxiliary variables + def get_writer(writer_map, name, factory, is_last): + """ + Return existing writer_map[name] or create via factory(). + In both cases set `.close_file = is_last`. + """ + initialization = False + if name not in writer_map: + writer_map[name] = factory() + initialization = True + w = writer_map[name] + w.close_file = is_last + + return w, initialization + output_writers = {} + compliance_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") + qs = Finding.all_objects.filter(scan_id=scan_id).order_by("uid").iterator() + for batch, is_last in batched(qs, DJANGO_FINDINGS_BATCH_SIZE): + fos = [FindingOutput.transform_api_finding(f, prowler_provider) for f in batch] - # 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", {})) + # Outputs + for mode, cfg in OUTPUT_FORMATS_MAPPING.items(): + cls = cfg["class"] + suffix = cfg["suffix"] + extra = cfg.get("kwargs", {}).copy() if mode == "html": - kwargs["provider"] = prowler_provider - kwargs["stats"] = scan_summary + extra.update(provider=prowler_provider, 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"], + writer, initialization = get_writer( + output_writers, + cls, + lambda cls=cls, fos=fos, suffix=suffix: cls( + findings=fos, + file_path=out_dir, + file_extension=suffix, from_cli=False, - ) - writer.close_file = is_last_batch - output_writers[writer_class] = writer + ), + is_last, + ) + if not initialization: + writer.transform(fos) + writer.batch_write_data_to_file(**extra) + writer._data.clear() - # Write the current batch using the writer - writer.batch_write_data_to_file(**kwargs) + # Compliance CSVs + for name in frameworks_avail: + compliance_obj = frameworks_bulk[name] - # TODO: Refactor the output classes to avoid this manual reset - writer._data = [] + klass = GenericCompliance + for condition, cls in COMPLIANCE_CLASS_MAP.get(provider_type, []): + if condition(name): + klass = cls + break - # Compress output files - output_directory = _compress_output_files(output_directory) + filename = f"{comp_dir}_{name}.csv" - # Save to configured storage - uploaded = _upload_to_s3(tenant_id, output_directory, scan_id) + writer, initialization = get_writer( + compliance_writers, + name, + lambda klass=klass, fos=fos: klass( + findings=fos, + compliance=compliance_obj, + file_path=filename, + from_cli=False, + ), + is_last, + ) + if not initialization: + writer.transform(fos, compliance_obj, name) + writer.batch_write_data_to_file() + writer._data.clear() - if uploaded: - # Remove the local files after upload + compressed = _compress_output_files(out_dir) + upload_uri = _upload_to_s3(tenant_id, compressed, scan_id) + + if upload_uri: try: - rmtree(Path(output_directory).parent, ignore_errors=True) - except FileNotFoundError as e: + rmtree(Path(compressed).parent, ignore_errors=True) + except Exception as e: logger.error(f"Error deleting output files: {e}") - - output_directory = uploaded - uploaded = True + final_location, did_upload = upload_uri, True else: - uploaded = False + final_location, did_upload = compressed, 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.all_objects.filter(id=scan_id).update(output_location=final_location) + logger.info(f"Scan outputs at {final_location}") + return {"upload": did_upload} diff --git a/api/src/backend/tasks/tests/test_export.py b/api/src/backend/tasks/tests/test_export.py new file mode 100644 index 0000000000..2aefbce9f9 --- /dev/null +++ b/api/src/backend/tasks/tests/test_export.py @@ -0,0 +1,142 @@ +import os +import zipfile +from unittest.mock import MagicMock, patch + +import pytest +from botocore.exceptions import ClientError +from tasks.jobs.export import ( + _compress_output_files, + _generate_output_directory, + _upload_to_s3, + get_s3_client, +) + + +@pytest.mark.django_db +class TestOutputs: + def test_compress_output_files_creates_zip(self, tmp_path): + output_dir = tmp_path / "output" + output_dir.mkdir() + file_path = output_dir / "result.csv" + file_path.write_text("data") + + zip_path = _compress_output_files(str(output_dir)) + + assert zip_path.endswith(".zip") + assert os.path.exists(zip_path) + with zipfile.ZipFile(zip_path, "r") as zipf: + assert "output/result.csv" in zipf.namelist() + + @patch("tasks.jobs.export.boto3.client") + @patch("tasks.jobs.export.settings") + def test_get_s3_client_success(self, mock_settings, mock_boto_client): + mock_settings.DJANGO_OUTPUT_S3_AWS_ACCESS_KEY_ID = "test" + mock_settings.DJANGO_OUTPUT_S3_AWS_SECRET_ACCESS_KEY = "test" + mock_settings.DJANGO_OUTPUT_S3_AWS_SESSION_TOKEN = "token" + mock_settings.DJANGO_OUTPUT_S3_AWS_DEFAULT_REGION = "eu-west-1" + + client_mock = MagicMock() + mock_boto_client.return_value = client_mock + + client = get_s3_client() + assert client is not None + client_mock.list_buckets.assert_called() + + @patch("tasks.jobs.export.boto3.client") + @patch("tasks.jobs.export.settings") + def test_get_s3_client_fallback(self, mock_settings, mock_boto_client): + mock_boto_client.side_effect = [ + ClientError({"Error": {"Code": "403"}}, "ListBuckets"), + MagicMock(), + ] + client = get_s3_client() + assert client is not None + + @patch("tasks.jobs.export.get_s3_client") + @patch("tasks.jobs.export.base") + def test_upload_to_s3_success(self, mock_base, mock_get_client, tmp_path): + mock_base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET = "test-bucket" + + zip_path = tmp_path / "outputs.zip" + zip_path.write_bytes(b"dummy") + + compliance_dir = tmp_path / "compliance" + compliance_dir.mkdir() + compliance_file = compliance_dir / "report.csv" + compliance_file.write_text("ok") + + client_mock = MagicMock() + mock_get_client.return_value = client_mock + + result = _upload_to_s3("tenant-id", str(zip_path), "scan-id") + + expected_uri = "s3://test-bucket/tenant-id/scan-id/outputs.zip" + assert result == expected_uri + + assert client_mock.upload_file.call_count == 2 + + @patch("tasks.jobs.export.get_s3_client") + @patch("tasks.jobs.export.base") + def test_upload_to_s3_missing_bucket(self, mock_base, mock_get_client): + mock_base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET = "" + result = _upload_to_s3("tenant", "/tmp/fake.zip", "scan") + assert result is None + + @patch("tasks.jobs.export.get_s3_client") + @patch("tasks.jobs.export.base") + def test_upload_to_s3_skips_non_files(self, mock_base, mock_get_client, tmp_path): + mock_base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET = "test-bucket" + zip_path = tmp_path / "results.zip" + zip_path.write_bytes(b"zip") + + compliance_dir = tmp_path / "compliance" + compliance_dir.mkdir() + (compliance_dir / "subdir").mkdir() + + client_mock = MagicMock() + mock_get_client.return_value = client_mock + + result = _upload_to_s3("tenant", str(zip_path), "scan") + + expected_uri = "s3://test-bucket/tenant/scan/results.zip" + assert result == expected_uri + + client_mock.upload_file.assert_called_once() + + @patch( + "tasks.jobs.export.get_s3_client", + side_effect=ClientError({"Error": {}}, "Upload"), + ) + @patch("tasks.jobs.export.base") + @patch("tasks.jobs.export.logger.error") + def test_upload_to_s3_failure_logs_error( + self, mock_logger, mock_base, mock_get_client, tmp_path + ): + mock_base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET = "bucket" + zip_path = tmp_path / "zipfile.zip" + zip_path.write_bytes(b"zip") + + compliance_dir = tmp_path / "compliance" + compliance_dir.mkdir() + (compliance_dir / "report.csv").write_text("csv") + + _upload_to_s3("tenant", str(zip_path), "scan") + mock_logger.assert_called() + + def test_generate_output_directory_creates_paths(self, tmp_path): + from prowler.config.config import output_file_timestamp + + base_dir = str(tmp_path) + tenant_id = "t1" + scan_id = "s1" + provider = "aws" + + 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 path.endswith(f"{provider}-{output_file_timestamp}") + assert compliance.endswith(f"{provider}-{output_file_timestamp}") diff --git a/api/src/backend/tasks/tests/test_tasks.py b/api/src/backend/tasks/tests/test_tasks.py new file mode 100644 index 0000000000..5f7c3fd94f --- /dev/null +++ b/api/src/backend/tasks/tests/test_tasks.py @@ -0,0 +1,415 @@ +import uuid +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from tasks.tasks import generate_outputs + + +@pytest.mark.django_db +class TestGenerateOutputs: + 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.tasks.ScanSummary.objects.filter") as mock_filter: + mock_filter.return_value.exists.return_value = False + + result = generate_outputs( + scan_id=self.scan_id, + provider_id=self.provider_id, + tenant_id=self.tenant_id, + ) + + assert result == {"upload": False} + mock_filter.assert_called_once_with(scan_id=self.scan_id) + + @patch("tasks.tasks.rmtree") + @patch("tasks.tasks._upload_to_s3") + @patch("tasks.tasks._compress_output_files") + @patch("tasks.tasks.get_compliance_frameworks") + @patch("tasks.tasks.Compliance.get_bulk") + @patch("tasks.tasks.initialize_prowler_provider") + @patch("tasks.tasks.Provider.objects.get") + @patch("tasks.tasks.ScanSummary.objects.filter") + @patch("tasks.tasks.Finding.all_objects.filter") + def test_generate_outputs_happy_path( + self, + mock_finding_filter, + mock_scan_summary_filter, + mock_provider_get, + mock_initialize_provider, + mock_compliance_get_bulk, + mock_get_available_frameworks, + mock_compress, + 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 + + prowler_provider = MagicMock() + mock_initialize_provider.return_value = prowler_provider + + mock_compliance_get_bulk.return_value = {"cis": MagicMock()} + mock_get_available_frameworks.return_value = ["cis"] + + dummy_finding = MagicMock(uid="f1") + mock_finding_filter.return_value.order_by.return_value.iterator.return_value = [ + [dummy_finding], + True, + ] + + mock_transformed_stats = {"some": "stats"} + with ( + patch( + "tasks.tasks.FindingOutput._transform_findings_stats", + return_value=mock_transformed_stats, + ), + patch( + "tasks.tasks.FindingOutput.transform_api_finding", + return_value={"transformed": "f1"}, + ), + patch( + "tasks.tasks.OUTPUT_FORMATS_MAPPING", + { + "json": { + "class": MagicMock(name="JSONWriter"), + "suffix": ".json", + "kwargs": {}, + } + }, + ), + patch( + "tasks.tasks.COMPLIANCE_CLASS_MAP", + {"aws": [(lambda x: True, MagicMock(name="CSVCompliance"))]}, + ), + patch( + "tasks.tasks._generate_output_directory", + return_value=("out-dir", "comp-dir"), + ), + patch("tasks.tasks.Scan.all_objects.filter") as mock_scan_update, + ): + mock_compress.return_value = "/tmp/zipped.zip" + mock_upload.return_value = "s3://bucket/zipped.zip" + + result = generate_outputs( + scan_id=self.scan_id, + provider_id=self.provider_id, + tenant_id=self.tenant_id, + ) + + assert result == {"upload": True} + mock_scan_update.return_value.update.assert_called_once_with( + output_location="s3://bucket/zipped.zip" + ) + mock_rmtree.assert_called_once_with( + Path("/tmp/zipped.zip").parent, ignore_errors=True + ) + + def test_generate_outputs_fails_upload(self): + with ( + patch("tasks.tasks.ScanSummary.objects.filter") as mock_filter, + patch("tasks.tasks.Provider.objects.get"), + patch("tasks.tasks.initialize_prowler_provider"), + patch("tasks.tasks.Compliance.get_bulk"), + patch("tasks.tasks.get_compliance_frameworks"), + patch("tasks.tasks.Finding.all_objects.filter") as mock_findings, + patch( + "tasks.tasks._generate_output_directory", return_value=("out", "comp") + ), + patch("tasks.tasks.FindingOutput._transform_findings_stats"), + patch("tasks.tasks.FindingOutput.transform_api_finding"), + patch( + "tasks.tasks.OUTPUT_FORMATS_MAPPING", + { + "json": { + "class": MagicMock(name="Writer"), + "suffix": ".json", + "kwargs": {}, + } + }, + ), + patch( + "tasks.tasks.COMPLIANCE_CLASS_MAP", + {"aws": [(lambda x: True, MagicMock())]}, + ), + patch("tasks.tasks._compress_output_files", return_value="/tmp/compressed"), + patch("tasks.tasks._upload_to_s3", return_value=None), + patch("tasks.tasks.Scan.all_objects.filter") as mock_scan_update, + ): + mock_filter.return_value.exists.return_value = True + mock_findings.return_value.order_by.return_value.iterator.return_value = [ + [MagicMock()], + True, + ] + + result = generate_outputs( + scan_id="scan", + provider_id="provider", + tenant_id=self.tenant_id, + ) + + assert result == {"upload": False} + mock_scan_update.return_value.update.assert_called_once() + + def test_generate_outputs_triggers_html_extra_update(self): + mock_finding_output = MagicMock() + mock_finding_output.compliance = {"cis": ["requirement-1", "requirement-2"]} + + with ( + patch("tasks.tasks.ScanSummary.objects.filter") as mock_filter, + patch("tasks.tasks.Provider.objects.get"), + patch("tasks.tasks.initialize_prowler_provider"), + patch("tasks.tasks.Compliance.get_bulk", return_value={"cis": MagicMock()}), + patch("tasks.tasks.get_compliance_frameworks", return_value=["cis"]), + patch("tasks.tasks.Finding.all_objects.filter") as mock_findings, + patch( + "tasks.tasks._generate_output_directory", return_value=("out", "comp") + ), + patch( + "tasks.tasks.FindingOutput._transform_findings_stats", + return_value={"some": "stats"}, + ), + patch( + "tasks.tasks.FindingOutput.transform_api_finding", + return_value=mock_finding_output, + ), + patch("tasks.tasks._compress_output_files", return_value="/tmp/compressed"), + patch("tasks.tasks._upload_to_s3", return_value="s3://bucket/f.zip"), + patch("tasks.tasks.Scan.all_objects.filter"), + ): + mock_filter.return_value.exists.return_value = True + mock_findings.return_value.order_by.return_value.iterator.return_value = [ + [MagicMock()], + True, + ] + + html_writer_mock = MagicMock() + with ( + patch( + "tasks.tasks.OUTPUT_FORMATS_MAPPING", + { + "html": { + "class": lambda *args, **kwargs: html_writer_mock, + "suffix": ".html", + "kwargs": {}, + } + }, + ), + patch( + "tasks.tasks.COMPLIANCE_CLASS_MAP", + {"aws": [(lambda x: True, MagicMock())]}, + ), + ): + generate_outputs( + scan_id=self.scan_id, + provider_id=self.provider_id, + tenant_id=self.tenant_id, + ) + html_writer_mock.batch_write_data_to_file.assert_called_once() + + def test_transform_called_only_on_second_batch(self): + raw1 = MagicMock() + raw2 = MagicMock() + + tf1 = MagicMock() + tf1.compliance = {} + tf2 = MagicMock() + tf2.compliance = {} + + writer_instances = [] + + class TrackingWriter: + def __init__(self, findings, file_path, file_extension, from_cli): + self.transform_called = 0 + self.batch_write_data_to_file = MagicMock() + self._data = [] + self.close_file = False + writer_instances.append(self) + + def transform(self, fos): + self.transform_called += 1 + + with ( + patch("tasks.tasks.ScanSummary.objects.filter") as mock_summary, + patch("tasks.tasks.Provider.objects.get"), + patch("tasks.tasks.initialize_prowler_provider"), + patch("tasks.tasks.Compliance.get_bulk"), + patch("tasks.tasks.get_compliance_frameworks", return_value=[]), + patch("tasks.tasks.FindingOutput._transform_findings_stats"), + patch( + "tasks.tasks.FindingOutput.transform_api_finding", + side_effect=[tf1, tf2], + ), + patch( + "tasks.tasks._generate_output_directory", + return_value=("outdir", "compdir"), + ), + patch("tasks.tasks._compress_output_files", return_value="outdir.zip"), + patch("tasks.tasks._upload_to_s3", return_value="s3://bucket/outdir.zip"), + patch("tasks.tasks.rmtree"), + patch("tasks.tasks.Scan.all_objects.filter"), + patch( + "tasks.tasks.batched", + return_value=[ + ([raw1], False), + ([raw2], True), + ], + ), + ): + mock_summary.return_value.exists.return_value = True + + with patch( + "tasks.tasks.OUTPUT_FORMATS_MAPPING", + { + "json": { + "class": TrackingWriter, + "suffix": ".json", + "kwargs": {}, + } + }, + ): + result = generate_outputs( + scan_id=self.scan_id, + provider_id=self.provider_id, + tenant_id=self.tenant_id, + ) + + assert result == {"upload": True} + assert len(writer_instances) == 1 + writer = writer_instances[0] + assert writer.transform_called == 1 + + def test_compliance_transform_called_on_second_batch(self): + raw1 = MagicMock() + raw2 = MagicMock() + compliance_obj = MagicMock() + writer_instances = [] + + class TrackingComplianceWriter: + def __init__(self, *args, **kwargs): + self.transform_calls = [] + self._data = [] + writer_instances.append(self) + + def transform(self, fos, comp_obj, name): + self.transform_calls.append((fos, comp_obj, name)) + + def batch_write_data_to_file(self): + pass + + two_batches = [ + ([raw1], False), + ([raw2], True), + ] + + with ( + patch("tasks.tasks.ScanSummary.objects.filter") as mock_summary, + patch( + "tasks.tasks.Provider.objects.get", + return_value=MagicMock(uid="UID", provider="aws"), + ), + patch("tasks.tasks.initialize_prowler_provider"), + patch( + "tasks.tasks.Compliance.get_bulk", return_value={"cis": compliance_obj} + ), + patch("tasks.tasks.get_compliance_frameworks", return_value=["cis"]), + patch( + "tasks.tasks._generate_output_directory", + return_value=("outdir", "compdir"), + ), + patch("tasks.tasks.FindingOutput._transform_findings_stats"), + patch( + "tasks.tasks.FindingOutput.transform_api_finding", + side_effect=lambda f, prov: f, + ), + patch("tasks.tasks._compress_output_files", return_value="outdir.zip"), + patch("tasks.tasks._upload_to_s3", return_value="s3://bucket/outdir.zip"), + patch("tasks.tasks.rmtree"), + patch( + "tasks.tasks.Scan.all_objects.filter", + return_value=MagicMock(update=lambda **kw: None), + ), + patch("tasks.tasks.batched", return_value=two_batches), + patch("tasks.tasks.OUTPUT_FORMATS_MAPPING", {}), + patch( + "tasks.tasks.COMPLIANCE_CLASS_MAP", + {"aws": [(lambda name: True, TrackingComplianceWriter)]}, + ), + ): + mock_summary.return_value.exists.return_value = True + + result = generate_outputs( + scan_id=self.scan_id, + provider_id=self.provider_id, + tenant_id=self.tenant_id, + ) + + assert len(writer_instances) == 1 + writer = writer_instances[0] + assert writer.transform_calls == [([raw2], compliance_obj, "cis")] + assert result == {"upload": True} + + def test_generate_outputs_logs_rmtree_exception(self, caplog): + mock_finding_output = MagicMock() + mock_finding_output.compliance = {"cis": ["requirement-1", "requirement-2"]} + + with ( + patch("tasks.tasks.ScanSummary.objects.filter") as mock_filter, + patch("tasks.tasks.Provider.objects.get"), + patch("tasks.tasks.initialize_prowler_provider"), + patch("tasks.tasks.Compliance.get_bulk", return_value={"cis": MagicMock()}), + patch("tasks.tasks.get_compliance_frameworks", return_value=["cis"]), + patch("tasks.tasks.Finding.all_objects.filter") as mock_findings, + patch( + "tasks.tasks._generate_output_directory", return_value=("out", "comp") + ), + patch( + "tasks.tasks.FindingOutput._transform_findings_stats", + return_value={"some": "stats"}, + ), + patch( + "tasks.tasks.FindingOutput.transform_api_finding", + return_value=mock_finding_output, + ), + patch("tasks.tasks._compress_output_files", return_value="/tmp/compressed"), + patch("tasks.tasks._upload_to_s3", return_value="s3://bucket/file.zip"), + patch("tasks.tasks.Scan.all_objects.filter"), + patch("tasks.tasks.rmtree", side_effect=Exception("Test deletion error")), + ): + mock_filter.return_value.exists.return_value = True + mock_findings.return_value.order_by.return_value.iterator.return_value = [ + [MagicMock()], + True, + ] + + with ( + patch( + "tasks.tasks.OUTPUT_FORMATS_MAPPING", + { + "json": { + "class": lambda *args, **kwargs: MagicMock(), + "suffix": ".json", + "kwargs": {}, + } + }, + ), + patch( + "tasks.tasks.COMPLIANCE_CLASS_MAP", + {"aws": [(lambda x: True, MagicMock())]}, + ), + ): + with caplog.at_level("ERROR"): + generate_outputs( + scan_id=self.scan_id, + provider_id=self.provider_id, + tenant_id=self.tenant_id, + ) + assert "Error deleting output files" in caplog.text diff --git a/prowler/lib/outputs/compliance/compliance_output.py b/prowler/lib/outputs/compliance/compliance_output.py index bb0e49d452..d3b855521e 100644 --- a/prowler/lib/outputs/compliance/compliance_output.py +++ b/prowler/lib/outputs/compliance/compliance_output.py @@ -31,14 +31,21 @@ class ComplianceOutput(Output): compliance: Compliance, file_path: str = None, file_extension: str = "", + from_cli: bool = True, ) -> None: + # TODO: This class needs to be refactored to use the Output class init, methods and properties self._data = [] + self.close_file = False + self.file_path = file_path self.file_descriptor = None + # This parameter is to avoid refactoring more code, the CLI does not write in batches, the API does + self._from_cli = from_cli if not file_extension and file_path: self._file_extension = "".join(Path(file_path).suffixes) if file_extension: self._file_extension = file_extension + self.file_path = f"{file_path}{self.file_extension}" if findings: # Get the compliance name of the model @@ -49,7 +56,7 @@ class ComplianceOutput(Output): ) self.transform(findings, compliance, compliance_name) if not self._file_descriptor and file_path: - self.create_file_descriptor(file_path) + self.create_file_descriptor(self.file_path) def batch_write_data_to_file(self) -> None: """ @@ -69,12 +76,14 @@ class ComplianceOutput(Output): fieldnames=[field.upper() for field in self._data[0].dict().keys()], delimiter=";", ) - csv_writer.writeheader() + if self._file_descriptor.tell() == 0: + csv_writer.writeheader() for finding in self._data: csv_writer.writerow( {k.upper(): v for k, v in finding.dict().items()} ) - self._file_descriptor.close() + if self.close_file or self._from_cli: + self._file_descriptor.close() except Exception as error: logger.error( f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}" diff --git a/tests/lib/outputs/compliance/aws_well_architected/aws_well_architected_test.py b/tests/lib/outputs/compliance/aws_well_architected/aws_well_architected_test.py index a53d4b5c30..104c127d6e 100644 --- a/tests/lib/outputs/compliance/aws_well_architected/aws_well_architected_test.py +++ b/tests/lib/outputs/compliance/aws_well_architected/aws_well_architected_test.py @@ -25,7 +25,9 @@ class TestAWSWellArchitected: ) ] - output = AWSWellArchitected(findings, AWS_WELL_ARCHITECTED) + output = AWSWellArchitected( + findings, AWS_WELL_ARCHITECTED, file_extension=".csv" + ) output_data = output.data[0] assert isinstance(output_data, AWSWellArchitectedModel) assert output_data.Provider == "aws"