# Copyright 2024 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

from debusine.artifacts import SigningKeyArtifact, WorkRequestDebugLogs
from debusine.artifacts.models import EmptyArtifactData, KeyPurpose
from debusine.client.debusine import Debusine
from debusine.signing.db.models import Key
from debusine.signing.gnupg import gpg_ephemeral_context
from debusine.signing.tasks.generate_key import GenerateKey
from debusine.test import TestCase


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

    def mock_debusine(self, task: GenerateKey) -> mock.MagicMock:
        """Create a Debusine mock and configure a task for it."""
        debusine_mock = mock.create_autospec(spec=Debusine)
        task.configure_server_access(debusine_mock)
        return debusine_mock

    def test_uefi(self) -> None:
        """Integration test: generate and store a real UEFI key."""
        service_private_key = PrivateKey.generate()
        task = GenerateKey(
            task_data={"purpose": "uefi", "description": "A UEFI key"},
            dynamic_task_data={},
        )
        debusine_mock = self.mock_debusine(task)
        task.work_request_id = 1
        task.workspace_name = "System"
        debug_log_files_directory = TemporaryDirectory(prefix="debusine-tests-")
        self.addCleanup(debug_log_files_directory.cleanup)
        task._debug_log_files_directory = debug_log_files_directory

        with override_settings(
            DEBUSINE_SIGNING_PRIVATE_KEYS=[service_private_key]
        ):
            self.assertTrue(task.execute())

        # 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 uploaded an artifact with the expected properties.
        debusine_mock.upload_artifact.assert_has_calls(
            [
                mock.call(
                    SigningKeyArtifact.create(
                        KeyPurpose.UEFI, key.fingerprint, key.public_key
                    ),
                    workspace="System",
                    work_request=1,
                ),
                mock.call(
                    WorkRequestDebugLogs(
                        category=WorkRequestDebugLogs._category,
                        data=EmptyArtifactData(),
                        files={
                            "cmd-output.log": (
                                Path(debug_log_files_directory.name)
                                / "cmd-output.log"
                            )
                        },
                    ),
                    workspace="System",
                    work_request=1,
                ),
            ]
        )

    def test_openpgp(self) -> None:
        """Integration test: generate and store a real OpenPGP key."""
        service_private_key = PrivateKey.generate()
        task = GenerateKey(
            task_data={"purpose": "openpgp", "description": "An OpenPGP key"},
            dynamic_task_data={},
        )
        debusine_mock = self.mock_debusine(task)
        task.work_request_id = 1
        task.workspace_name = "System"
        debug_log_files_directory = TemporaryDirectory(prefix="debusine-tests-")
        self.addCleanup(debug_log_files_directory.cleanup)
        task._debug_log_files_directory = debug_log_files_directory

        with override_settings(
            DEBUSINE_SIGNING_PRIVATE_KEYS=[service_private_key]
        ):
            self.assertTrue(task.execute())

        # 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)
        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 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 uploaded an artifact with the expected properties.
        debusine_mock.upload_artifact.assert_has_calls(
            [
                mock.call(
                    SigningKeyArtifact.create(
                        KeyPurpose.OPENPGP, key.fingerprint, key.public_key
                    ),
                    workspace="System",
                    work_request=1,
                ),
                mock.call(
                    WorkRequestDebugLogs(
                        category=WorkRequestDebugLogs._category,
                        data=EmptyArtifactData(),
                        files={
                            "cmd-output.log": (
                                Path(debug_log_files_directory.name)
                                / "cmd-output.log"
                            )
                        },
                    ),
                    workspace="System",
                    work_request=1,
                ),
            ]
        )

    def test_generate_failure(self) -> None:
        """If key generation fails, the task does not upload a signing key."""
        task = GenerateKey(
            task_data={"purpose": "uefi", "description": "A UEFI key"},
            dynamic_task_data={},
        )
        debusine_mock = self.mock_debusine(task)
        task.work_request_id = 1
        task.workspace_name = "System"
        debug_log_files_directory = TemporaryDirectory(prefix="debusine-tests-")
        self.addCleanup(debug_log_files_directory.cleanup)
        task._debug_log_files_directory = debug_log_files_directory

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

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

        # Only a debug logs artifact was uploaded, not a signing key artifact.
        debusine_mock.upload_artifact.assert_called_once_with(
            WorkRequestDebugLogs(
                category=WorkRequestDebugLogs._category,
                data=EmptyArtifactData(),
                files={
                    "cmd-output.log": (
                        Path(debug_log_files_directory.name) / "cmd-output.log"
                    ),
                    "execution.log": (
                        Path(debug_log_files_directory.name) / "execution.log"
                    ),
                },
            ),
            workspace="System",
            work_request=1,
        )

    def test_label(self) -> None:
        """Test get_label."""
        task = GenerateKey(
            task_data={"purpose": "uefi", "description": "A UEFI key"},
            dynamic_task_data={},
        )
        self.assertEqual(task.get_label(), "generate uefi key")
