# Copyright © The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

"""Unit tests for the GenerateKey task."""

import base64
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest import mock

from cryptography.x509 import load_pem_x509_certificate
from django.test import TransactionTestCase, override_settings
from nacl.public import PrivateKey, SealedBox

import debusine.worker.tags as wtags
from debusine.artifacts import WorkRequestDebugLogs
from debusine.artifacts.models import EmptyArtifactData, WorkRequestResults
from debusine.assets import AssetCategory, KeyPurpose, SigningKeyData
from debusine.signing.db.models import Key
from debusine.signing.gnupg import gpg_ephemeral_context
from debusine.signing.tasks import GenerateKey
from debusine.signing.tasks.tests.helpers import SigningTaskTestCase
from debusine.test import TestCase


class GenerateKeyTests(
    SigningTaskTestCase[GenerateKey], TestCase, TransactionTestCase
):
    """Test the :py:class:`GenerateKey` task."""

    SAMPLE_TASK_DATA = {"purpose": "uefi", "description": "A UEFI key"}

    def test_compute_scheduler_tags(self) -> None:
        self.assertComputesSchedulerTags()

    def test_get_subject(self) -> None:
        self.assertIsNone(self.get_subject())

    def test_get_configuration_context(self) -> None:
        self.assertIsNone(self.get_configuration_context())

    def test_get_provided_worker_tags(self) -> None:
        self.assertCountEqual(
            self.task.get_provided_worker_tags(),
            [wtags.TASK_PREFIX + "signing:generatekey:version:1"],
        )

    def test_uefi(self) -> None:
        """Integration test: generate and store a real UEFI key."""
        service_private_key = PrivateKey.generate()
        debusine_mock = self.mock_debusine(self.task)

        with override_settings(
            DEBUSINE_SIGNING_PRIVATE_KEYS=[service_private_key]
        ):
            self.assertEqual(self.task.execute(), WorkRequestResults.SUCCESS)

        # The private key was encrypted using the configured service-wide
        # private key, and the public key has the correct subject.
        key = Key.objects.get(purpose=Key.Purpose.UEFI)
        self.assertEqual(
            load_pem_x509_certificate(
                bytes(key.public_key)
            ).subject.rfc4514_string(),
            "CN=A UEFI key",
        )
        self.assertEqual(key.private_key["storage"], "nacl")
        self.assertEqual(
            base64.b64decode(key.private_key["public_key"].encode()),
            bytes(service_private_key.public_key),
        )
        SealedBox(service_private_key).decrypt(
            base64.b64decode(key.private_key["encrypted"].encode())
        )

        # The task created an asset with the expected properties.
        debusine_mock.asset_create.assert_called_once_with(
            category=AssetCategory.SIGNING_KEY,
            data=SigningKeyData(
                purpose=KeyPurpose.UEFI,
                fingerprint=key.fingerprint,
                public_key=base64.b64encode(key.public_key).decode(),
                description="A UEFI key",
            ),
            workspace=self.task.workspace_name,
            work_request=self.task.work_request_id,
        )

    def test_openpgp_upload(self) -> None:
        """Integration test: generate and store a real OpenPGP upload key."""
        self.task_data = {
            "purpose": "openpgp",
            "description": "An OpenPGP upload key",
        }
        service_private_key = PrivateKey.generate()
        debusine_mock = self.mock_debusine(self.task)

        with override_settings(
            DEBUSINE_SIGNING_PRIVATE_KEYS=[service_private_key]
        ):
            self.assertEqual(self.task.execute(), WorkRequestResults.SUCCESS)

        # The private key was encrypted using the configured service-wide
        # private key, and the public key has the correct UID.
        key = Key.objects.get(purpose=Key.Purpose.OPENPGP_UPLOAD)
        with (
            TemporaryDirectory(prefix="debusine-tests-") as tmp,
            gpg_ephemeral_context(Path(tmp)) as ctx,
        ):
            result = ctx.key_import(bytes(key.public_key))
            assert getattr(result, "imported", 0) == 1
            gpg_key = ctx.get_key(result.imports[0].fpr)
            self.assertEqual(gpg_key.uids[0].name, "An OpenPGP upload key")
        self.assertEqual(key.private_key["storage"], "nacl")
        self.assertEqual(
            base64.b64decode(key.private_key["public_key"].encode()),
            bytes(service_private_key.public_key),
        )
        SealedBox(service_private_key).decrypt(
            base64.b64decode(key.private_key["encrypted"].encode())
        )

        # The task created a signing asset with the expected properties.
        debusine_mock.asset_create.assert_called_once_with(
            category=AssetCategory.SIGNING_KEY,
            data=SigningKeyData(
                purpose=KeyPurpose.OPENPGP_UPLOAD,
                fingerprint=key.fingerprint,
                public_key=base64.b64encode(key.public_key).decode(),
                description="An OpenPGP upload key",
            ),
            workspace=self.task.workspace_name,
            work_request=self.task.work_request_id,
        )

    def test_openpgp_repository(self) -> None:
        """Integration test: generate and store a real OpenPGP repo key."""
        self.task_data = {
            "purpose": "openpgp-repository",
            "description": "An OpenPGP repository key",
        }
        service_private_key = PrivateKey.generate()
        debusine_mock = self.mock_debusine(self.task)

        with override_settings(
            DEBUSINE_SIGNING_PRIVATE_KEYS=[service_private_key]
        ):
            self.assertEqual(self.task.execute(), WorkRequestResults.SUCCESS)

        # The private key was encrypted using the configured service-wide
        # private key, and the public key has the correct UID.
        key = Key.objects.get(purpose=Key.Purpose.OPENPGP_REPOSITORY)
        with (
            TemporaryDirectory(prefix="debusine-tests-") as tmp,
            gpg_ephemeral_context(Path(tmp)) as ctx,
        ):
            result = ctx.key_import(bytes(key.public_key))
            assert getattr(result, "imported", 0) == 1
            gpg_key = ctx.get_key(result.imports[0].fpr)
            self.assertEqual(gpg_key.uids[0].name, "An OpenPGP repository key")
        self.assertEqual(key.private_key["storage"], "nacl")
        self.assertEqual(
            base64.b64decode(key.private_key["public_key"].encode()),
            bytes(service_private_key.public_key),
        )
        SealedBox(service_private_key).decrypt(
            base64.b64decode(key.private_key["encrypted"].encode())
        )

        # The task created a signing asset with the expected properties.
        debusine_mock.asset_create.assert_called_once_with(
            category=AssetCategory.SIGNING_KEY,
            data=SigningKeyData(
                purpose=KeyPurpose.OPENPGP_REPOSITORY,
                fingerprint=key.fingerprint,
                public_key=base64.b64encode(key.public_key).decode(),
                description="An OpenPGP repository key",
            ),
            workspace=self.task.workspace_name,
            work_request=self.task.work_request_id,
        )

    def test_generate_failure(self) -> None:
        """If key generation fails, the task does not upload a signing key."""
        debusine_mock = self.mock_debusine(self.task)

        with mock.patch.object(
            Key.objects,
            "generate",
            side_effect=Exception("Key generation failed"),
        ):
            self.assertEqual(self.task.execute(), WorkRequestResults.FAILURE)

        # No key was generated.
        self.assertQuerySetEqual(Key.objects.all(), [])

        # No asset was created
        debusine_mock.asset_create.assert_not_called()

        # Debug logs artifact was uploaded
        debusine_mock.upload_artifact.assert_called_once_with(
            WorkRequestDebugLogs(
                category=WorkRequestDebugLogs._category,
                data=EmptyArtifactData(),
                files={
                    "stages.log": self.task.log_directory / "stages.log",
                    "cmd-output.log": (
                        self.task.log_directory / "cmd-output.log"
                    ),
                    "execution.log": (
                        self.task.log_directory / "execution.log"
                    ),
                },
                content_types={
                    "stages.log": "text/plain; charset=utf-8",
                    "cmd-output.log": "text/plain; charset=utf-8",
                    "execution.log": "text/plain; charset=utf-8",
                },
            ),
            workspace=self.task.workspace_name,
            work_request=self.task.work_request_id,
        )
