diff --git a/api/src/backend/tasks/jobs/export.py b/api/src/backend/tasks/jobs/export.py index 7b4ead1778..f9232e1ed0 100644 --- a/api/src/backend/tasks/jobs/export.py +++ b/api/src/backend/tasks/jobs/export.py @@ -8,11 +8,12 @@ from botocore.exceptions import ClientError, NoCredentialsError, ParamValidation from celery.utils.log import get_task_logger from django.conf import settings +from api.db_utils import rls_transaction +from api.models import Scan from prowler.config.config import ( csv_file_suffix, html_file_suffix, json_ocsf_file_suffix, - output_file_timestamp, ) from prowler.lib.outputs.compliance.aws_well_architected.aws_well_architected import ( AWSWellArchitected, @@ -248,15 +249,19 @@ def _generate_output_directory( # Sanitize the prowler provider name to ensure it is a valid directory name prowler_provider_sanitized = re.sub(r"[^\w\-]", "-", prowler_provider) + with rls_transaction(tenant_id): + started_at = Scan.objects.get(id=scan_id).started_at + + timestamp = started_at.strftime("%Y%m%d%H%M%S") path = ( f"{output_directory}/{tenant_id}/{scan_id}/prowler-output-" - f"{prowler_provider_sanitized}-{output_file_timestamp}" + f"{prowler_provider_sanitized}-{timestamp}" ) os.makedirs("/".join(path.split("/")[:-1]), exist_ok=True) compliance_path = ( f"{output_directory}/{tenant_id}/{scan_id}/compliance/prowler-output-" - f"{prowler_provider_sanitized}-{output_file_timestamp}" + f"{prowler_provider_sanitized}-{timestamp}" ) os.makedirs("/".join(compliance_path.split("/")[:-1]), exist_ok=True) diff --git a/api/src/backend/tasks/tests/test_export.py b/api/src/backend/tasks/tests/test_export.py index f113d22f17..c10f20774b 100644 --- a/api/src/backend/tasks/tests/test_export.py +++ b/api/src/backend/tasks/tests/test_export.py @@ -1,5 +1,7 @@ import os +import uuid import zipfile +from datetime import datetime from pathlib import Path from unittest.mock import MagicMock, patch @@ -127,14 +129,26 @@ class TestOutputs: _upload_to_s3("tenant", str(zip_path), "scan") mock_logger.assert_called() - def test_generate_output_directory_creates_paths(self, tmpdir): - from prowler.config.config import output_file_timestamp + @patch("tasks.jobs.export.rls_transaction") + @patch("tasks.jobs.export.Scan") + def test_generate_output_directory_creates_paths( + self, mock_scan, mock_rls_transaction, tmpdir + ): + # Mock the scan object with a started_at timestamp + mock_scan_instance = MagicMock() + mock_scan_instance.started_at = datetime(2023, 6, 15, 10, 30, 45) + mock_scan.objects.get.return_value = mock_scan_instance + + # Mock rls_transaction as a context manager + mock_rls_transaction.return_value.__enter__ = MagicMock() + mock_rls_transaction.return_value.__exit__ = MagicMock(return_value=False) base_tmp = Path(str(tmpdir.mkdir("generate_output"))) base_dir = str(base_tmp) - tenant_id = "t1" - scan_id = "s1" + tenant_id = str(uuid.uuid4()) + scan_id = str(uuid.uuid4()) provider = "aws" + expected_timestamp = "20230615103045" path, compliance = _generate_output_directory( base_dir, provider, tenant_id, scan_id @@ -143,17 +157,29 @@ class TestOutputs: 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}") + assert path.endswith(f"{provider}-{expected_timestamp}") + assert compliance.endswith(f"{provider}-{expected_timestamp}") - def test_generate_output_directory_invalid_character(self, tmpdir): - from prowler.config.config import output_file_timestamp + @patch("tasks.jobs.export.rls_transaction") + @patch("tasks.jobs.export.Scan") + def test_generate_output_directory_invalid_character( + self, mock_scan, mock_rls_transaction, tmpdir + ): + # Mock the scan object with a started_at timestamp + mock_scan_instance = MagicMock() + mock_scan_instance.started_at = datetime(2023, 6, 15, 10, 30, 45) + mock_scan.objects.get.return_value = mock_scan_instance + + # Mock rls_transaction as a context manager + mock_rls_transaction.return_value.__enter__ = MagicMock() + mock_rls_transaction.return_value.__exit__ = MagicMock(return_value=False) base_tmp = Path(str(tmpdir.mkdir("generate_output"))) base_dir = str(base_tmp) - tenant_id = "t1" - scan_id = "s1" + tenant_id = str(uuid.uuid4()) + scan_id = str(uuid.uuid4()) provider = "aws/test@check" + expected_timestamp = "20230615103045" path, compliance = _generate_output_directory( base_dir, provider, tenant_id, scan_id @@ -162,5 +188,5 @@ class TestOutputs: assert os.path.isdir(os.path.dirname(path)) assert os.path.isdir(os.path.dirname(compliance)) - assert path.endswith(f"aws-test-check-{output_file_timestamp}") - assert compliance.endswith(f"aws-test-check-{output_file_timestamp}") + assert path.endswith(f"aws-test-check-{expected_timestamp}") + assert compliance.endswith(f"aws-test-check-{expected_timestamp}")