prowler/api/src/backend/tasks/tests/test_export.py

import os
import zipfile
from pathlib import Path
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, tmpdir):
        base_tmp = Path(str(tmpdir.mkdir("compress_output")))
        output_dir = base_tmp / "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, tmpdir):
        mock_base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET = "test-bucket"

        base_tmp = Path(str(tmpdir.mkdir("upload_success")))
        zip_path = base_tmp / "outputs.zip"
        zip_path.write_bytes(b"dummy")

        compliance_dir = base_tmp / "compliance"
        compliance_dir.mkdir()
        (compliance_dir / "report.csv").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, tmpdir):
        mock_base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET = "test-bucket"
        base_tmp = Path(str(tmpdir.mkdir("upload_skips_non_files")))

        zip_path = base_tmp / "results.zip"
        zip_path.write_bytes(b"zip")

        compliance_dir = base_tmp / "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, tmpdir
    ):
        mock_base.DJANGO_OUTPUT_S3_AWS_OUTPUT_BUCKET = "bucket"

        base_tmp = Path(str(tmpdir.mkdir("upload_failure_logs")))
        zip_path = base_tmp / "zipfile.zip"
        zip_path.write_bytes(b"zip")

        compliance_dir = base_tmp / "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, tmpdir):
        from prowler.config.config import output_file_timestamp

        base_tmp = Path(str(tmpdir.mkdir("generate_output")))
        base_dir = str(base_tmp)
        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}")