# 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 Sign task."""

import os
import shutil
import subprocess
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Any
from unittest import mock, skipUnless

from django.test import TestCase, TransactionTestCase, override_settings
from django.utils import timezone

from nacl.public import PrivateKey

try:
    import pydantic.v1 as pydantic
except ImportError:
    import pydantic as pydantic  # type: ignore

from debusine.artifacts import (
    LocalArtifact,
    SigningInputArtifact,
    SigningKeyArtifact,
    SigningOutputArtifact,
    WorkRequestDebugLogs,
)
from debusine.artifacts.models import (
    DebusineSigningOutput,
    EmptyArtifactData,
    SigningResult,
)
from debusine.client.debusine import Debusine
from debusine.client.models import (
    ArtifactResponse,
    FileResponse,
    FilesResponseType,
    RemoteArtifact,
)
from debusine.signing.models import Key
from debusine.signing.tasks.models import KeyPurpose, SignDynamicData
from debusine.signing.tasks.sign import Sign
from debusine.tasks.models import WorkerType
from debusine.tasks.tests.helper_mixin import FakeTaskDatabase
from debusine.test import TestHelpersMixin
from debusine.utils import calculate_hash, is_command_available


# Bare minimum that sbsigntool will recognize as a valid PE/COFF image.
# This is unlikely to work in any other context.
minimal_pecoff_image = (
    # DOS header:
    # DOS magic number
    b"MZ"
    # DOS header fields, ignored by sbsigntool
    + (b"\0" * 58)
    # Address of next header
    + b"\x80\0\0\0"
    # DOS message, ignored by sbsigntool
    + (b"\0" * 64)
    # PE header:
    # NT signature
    + b"PE\0\0"
    # AMD64 magic
    + b"\x64\x86"
    # Number of sections
    + b"\0\0"
    # PE header fields, ignored by sbsigntool
    + (b"\0" * 12)
    # Optional header size: 752
    + (b"\xf0\x02")
    # Flags, ignored by sbsigntool
    + b"\0\0"
    # Optional header:
    # PE32+ magic
    + b"\x0b\x02"
    # Optional header standard fields, ignored by sbsigntool
    + (b"\0" * 22)
    # Image base and section alignment, ignored by sbsigntool
    + (b"\0" * 12)
    # File alignment: 512
    + b"\0\2\0\0"
    # Optional header NT-specific fields, ignored by sbsigntool
    + (b"\0" * 20)
    # Size of headers: 1024
    + b"\0\4\0\0"
    # More optional header NT-specific fields, ignored by sbsigntool
    + (b"\0" * 48)
    # Space for five data directories (the fifth is the certificate table)
    + (b"\0" * 640)
    # Padding for alignment
    + (b"\0" * 120)
)


example_url = pydantic.parse_obj_as(pydantic.AnyUrl, "http://example.org/")


class SignTestMixin:
    """Helpers for signing tests."""

    def mock_debusine(self, task: Sign) -> 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 make_files_response(self, files: dict[str, Path]) -> FilesResponseType:
        """Make a fake files response."""
        return FilesResponseType(
            {
                name: FileResponse(
                    size=path.stat().st_size,
                    checksums={"sha256": calculate_hash(path, "sha256").hex()},
                    type="file",
                    url=example_url,
                )
                for name, path in files.items()
            }
        )

    def make_artifact_response(
        self, artifact: LocalArtifact[Any], artifact_id: int
    ) -> ArtifactResponse:
        """Make a fake artifact response."""
        return ArtifactResponse(
            id=artifact_id,
            workspace="System",
            category=artifact.category,
            created_at=timezone.now(),
            data=artifact.data.dict(),
            download_tar_gz_url=example_url,
            files_to_upload=[],
            files=self.make_files_response(artifact.files),
        )


class SignTests(SignTestMixin, TestHelpersMixin, TestCase):
    """Unit tests for the :py:class:`Sign` task."""

    def test_analyze_worker(self) -> None:
        """Test the analyze_worker() method."""
        self.mock_is_command_available({"sbsign": True})
        task = Sign(
            task_data={"purpose": KeyPurpose.UEFI, "unsigned": 2, "key": 1}
        )
        metadata = task.analyze_worker()
        self.assertEqual(metadata["signing:sign:available:uefi"], True)

    def test_analyze_worker_sbsign_not_available(self) -> None:
        """analyze_worker() handles sbsign not being available."""
        self.mock_is_command_available({"sbsign": False})
        task = Sign(
            task_data={"purpose": KeyPurpose.UEFI, "unsigned": 2, "key": 1}
        )
        metadata = task.analyze_worker()
        self.assertEqual(metadata["signing:sign:available:uefi"], False)

    def test_can_run_on_uefi(self) -> None:
        """can_run_on returns True for UEFI if sbsign is available."""
        task = Sign(
            task_data={"purpose": KeyPurpose.UEFI, "unsigned": 2, "key": 1}
        )
        self.assertTrue(
            task.can_run_on(
                {
                    "system:worker_type": WorkerType.SIGNING,
                    "signing:sign:available:uefi": True,
                    "signing:sign:version": task.TASK_VERSION,
                }
            )
        )

    def test_can_run_on_mismatched_task_version(self) -> None:
        """can_run_on returns False for mismatched task versions."""
        task = Sign(
            task_data={"purpose": KeyPurpose.UEFI, "unsigned": 2, "key": 1}
        )
        self.assertFalse(
            task.can_run_on(
                {
                    "system:worker_type": WorkerType.SIGNING,
                    "signing:sign:available:uefi": True,
                    "signing:sign:version": task.TASK_VERSION + 1,
                }
            )
        )

    def test_can_run_on_uefi_missing_tool(self) -> None:
        """can_run_on returns False for UEFI if sbsign is not available."""
        task = Sign(
            task_data={"purpose": KeyPurpose.UEFI, "unsigned": 2, "key": 1}
        )
        self.assertFalse(
            task.can_run_on(
                {
                    "system:worker_type": WorkerType.SIGNING,
                    "signing:sign:available:uefi": False,
                    "signing:sign:version": task.TASK_VERSION,
                }
            )
        )

    def test_compute_dynamic_data(self) -> None:
        """Dynamic data receives relevant artifact IDs."""
        key_lookup = "bookworm@debian:suite-signing-keys/key:uefi"
        task_db = FakeTaskDatabase(
            single_lookups={
                # unsigned
                (2, None): 2,
                # key
                (key_lookup, None): 1,
            }
        )

        task = Sign(
            task_data={
                "purpose": KeyPurpose.UEFI,
                "unsigned": 2,
                "key": key_lookup,
            }
        )
        self.assertEqual(
            task.compute_dynamic_data(task_db),
            SignDynamicData(unsigned_id=2, key_id=1),
        )

    def test_fetch_input_missing_trusted_cert(self) -> None:
        """fetch_input requires trusted_certs to match configuration."""
        trusted_cert = "0" * 64
        signing_input = SigningInputArtifact.create(
            [], self.create_temporary_directory(), trusted_certs=[trusted_cert]
        )

        def download_artifact(
            artifact_id: int, destination: Path, **kwargs: Any  # noqa: U100
        ) -> ArtifactResponse:
            return self.make_artifact_response(signing_input, artifact_id)

        task = Sign(
            task_data={"purpose": KeyPurpose.UEFI, "unsigned": 2, "key": 1},
            dynamic_task_data={"unsigned_id": 2, "key_id": 1},
        )
        debusine_mock = self.mock_debusine(task)
        debusine_mock.download_artifact.side_effect = download_artifact
        debug_log_files_directory = TemporaryDirectory(prefix="debusine-tests-")
        self.addCleanup(debug_log_files_directory.cleanup)
        task._debug_log_files_directory = debug_log_files_directory
        destination = self.create_temporary_directory()

        with override_settings(DEBUSINE_SIGNING_TRUSTED_CERTS=[]):
            self.assertFalse(task.fetch_input(destination))

        self.assertEqual(
            Path(debug_log_files_directory.name, "fetch_input.log").read_text(),
            f"This installation cannot sign objects that trust certificate "
            f"'{trusted_cert}'\n",
        )

    def test_fetch_input_trusted_cert(self) -> None:
        """fetch_input accepts trusted_certs if they match configuration."""
        trusted_cert = "0" * 64
        fingerprint = "1" * 64
        key = Key.objects.create(
            purpose=Key.Purpose.UEFI,
            fingerprint=fingerprint,
            private_key={},
            public_key=b"",
        )
        signing_key = SigningKeyArtifact.create(
            KeyPurpose.UEFI, fingerprint=fingerprint, public_key=b""
        )
        signing_key_id = 1
        signing_input = SigningInputArtifact.create(
            [], self.create_temporary_directory(), trusted_certs=[trusted_cert]
        )
        signing_input_id = 2
        responses = {
            response.id: response
            for response in (
                self.make_artifact_response(signing_key, signing_key_id),
                self.make_artifact_response(signing_input, signing_input_id),
            )
        }

        def download_artifact(
            artifact_id: int, destination: Path, **kwargs: Any  # noqa: U100
        ) -> ArtifactResponse:
            return responses[artifact_id]

        task = Sign(
            task_data={
                "purpose": KeyPurpose.UEFI,
                "unsigned": signing_input_id,
                "key": signing_key_id,
            },
            dynamic_task_data={
                "unsigned_id": signing_input_id,
                "key_id": signing_key_id,
            },
        )
        debusine_mock = self.mock_debusine(task)
        debusine_mock.download_artifact.side_effect = download_artifact
        destination = self.create_temporary_directory()

        with override_settings(DEBUSINE_SIGNING_TRUSTED_CERTS=[trusted_cert]):
            self.assertTrue(task.fetch_input(destination))

        self.assertEqual(task._files, responses[2].files)
        self.assertEqual(task._key, key)

    def test_fetch_input_missing_key(self) -> None:
        """fetch_input requires the signing key to exist in the database."""
        fingerprint = "0" * 64
        signing_key = SigningKeyArtifact.create(
            KeyPurpose.UEFI, fingerprint=fingerprint, public_key=b""
        )
        signing_key_id = 1
        signing_input = SigningInputArtifact.create(
            [], self.create_temporary_directory()
        )
        signing_input_id = 2
        responses = {
            response.id: response
            for response in (
                self.make_artifact_response(signing_key, signing_key_id),
                self.make_artifact_response(signing_input, signing_input_id),
            )
        }

        def download_artifact(
            artifact_id: int, destination: Path, **kwargs: Any  # noqa: U100
        ) -> ArtifactResponse:
            return responses[artifact_id]

        task = Sign(
            task_data={
                "purpose": KeyPurpose.UEFI,
                "unsigned": signing_input_id,
                "key": signing_key_id,
            },
            dynamic_task_data={
                "unsigned_id": signing_input_id,
                "key_id": signing_key_id,
            },
        )
        debusine_mock = self.mock_debusine(task)
        debusine_mock.download_artifact.side_effect = download_artifact
        debug_log_files_directory = TemporaryDirectory(prefix="debusine-tests-")
        self.addCleanup(debug_log_files_directory.cleanup)
        task._debug_log_files_directory = debug_log_files_directory
        destination = self.create_temporary_directory()

        self.assertFalse(task.fetch_input(destination))

        self.assertEqual(
            Path(debug_log_files_directory.name, "fetch_input.log").read_text(),
            f"Signing key uefi:{fingerprint} does not exist\n",
        )

    def test_run_error(self) -> None:
        """run() records an error result if signing fails."""
        key = Key.objects.create(
            purpose=Key.Purpose.UEFI,
            fingerprint="0" * 64,
            private_key={},
            public_key=b"",
        )
        download_directory = self.create_temporary_directory()
        execute_directory = self.create_temporary_directory()
        (image := download_directory / "image").touch()
        task = Sign(
            task_data={"purpose": KeyPurpose.UEFI, "unsigned": 2, "key": 1}
        )
        task.work_request_id = 1
        task._files = self.make_files_response({"image": image})
        task._key = key
        task.prepare_to_run(download_directory, execute_directory)

        with mock.patch.object(key, "sign", side_effect=Exception("Boom")):
            self.assertTrue(task.run(execute_directory))

        self.assertEqual(
            task._results, [SigningResult(file="image", error_message="Boom")]
        )

    def test_run_success(self) -> None:
        """run() records a successful result if signing succeeds."""
        key = Key.objects.create(
            purpose=Key.Purpose.UEFI,
            fingerprint="0" * 64,
            private_key={},
            public_key=b"",
        )
        download_directory = self.create_temporary_directory()
        execute_directory = self.create_temporary_directory()
        (image := download_directory / "image").touch()
        task = Sign(
            task_data={"purpose": KeyPurpose.UEFI, "unsigned": 2, "key": 1}
        )
        task.work_request_id = 1
        task._files = self.make_files_response({"image": image})
        task._key = key
        task.prepare_to_run(download_directory, execute_directory)

        with mock.patch.object(key, "sign") as mock_sign:
            self.assertTrue(task.run(execute_directory))

        self.assertEqual(
            task._results,
            [SigningResult(file="image", output_file="image.sig")],
        )
        mock_sign.assert_called_once_with(
            execute_directory / "input" / "image",
            execute_directory / "output" / "image.sig",
            1,
            log_file=mock.ANY,
        )


class SignIntegrationTests(
    SignTestMixin, TestHelpersMixin, TransactionTestCase
):
    """Integration tests for the :py:class:`Sign` task."""

    @skipUnless(is_command_available("sbsign"), "requires sbsign")
    def test_uefi(self) -> None:
        """Integration test: sign using a real UEFI key."""
        service_private_key = PrivateKey.generate()
        with (
            override_settings(
                DEBUSINE_SIGNING_PRIVATE_KEYS=[service_private_key]
            ),
            open(os.devnull, "wb") as log_file,
        ):
            key = Key.objects.generate(
                Key.Purpose.UEFI, "A UEFI key", 1, log_file
            )
        signing_key = SigningKeyArtifact.create(
            KeyPurpose.UEFI,
            fingerprint=key.fingerprint,
            public_key=key.public_key,
        )
        signing_key_id = 1
        temp_path = self.create_temporary_directory()
        (minimal_pecoff_path := temp_path / "image").write_bytes(
            minimal_pecoff_image
        )
        signing_input = SigningInputArtifact.create(
            [minimal_pecoff_path], temp_path
        )
        signing_input_id = 2
        responses = {
            response.id: response
            for response in (
                self.make_artifact_response(signing_key, signing_key_id),
                self.make_artifact_response(signing_input, signing_input_id),
            )
        }
        uploaded_paths: list[Path] = []
        output_path = self.create_temporary_directory()
        signing_output_id = 3

        def download_artifact(
            artifact_id: int, destination: Path, **kwargs: Any  # noqa: U100
        ) -> ArtifactResponse:
            if artifact_id == signing_input_id:
                shutil.copy(minimal_pecoff_path, destination)
            return responses[artifact_id]

        def upload_artifact(
            local_artifact: LocalArtifact[Any], **kwargs: Any
        ) -> RemoteArtifact:
            nonlocal uploaded_paths
            for path in local_artifact.files.values():
                uploaded_paths.append(path)
                shutil.copy(path, output_path)
            return RemoteArtifact(id=signing_output_id, workspace="System")

        task = Sign(
            task_data={
                "purpose": KeyPurpose.UEFI,
                "unsigned": signing_input_id,
                "key": signing_key_id,
            },
            dynamic_task_data={
                "unsigned_id": signing_input_id,
                "key_id": signing_key_id,
            },
        )
        debusine_mock = self.mock_debusine(task)
        debusine_mock.download_artifact.side_effect = download_artifact
        debusine_mock.upload_artifact.side_effect = upload_artifact
        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 task uploaded an artifact with the expected properties.
        self.assertEqual(len(uploaded_paths), 2)
        debusine_mock.upload_artifact.assert_has_calls(
            [
                mock.call(
                    SigningOutputArtifact(
                        category=SigningOutputArtifact._category,
                        files={"image.sig": uploaded_paths[0]},
                        data=DebusineSigningOutput(
                            purpose=KeyPurpose.UEFI,
                            fingerprint=key.fingerprint,
                            results=[
                                SigningResult(
                                    file="image", output_file="image.sig"
                                )
                            ],
                        ),
                    ),
                    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,
                ),
            ]
        )
        debusine_mock.relation_create.assert_has_calls(
            [
                mock.call(signing_output_id, signing_input_id, "relates-to"),
                mock.call(signing_output_id, signing_key_id, "relates-to"),
            ]
        )

        # The file was properly signed.
        (certificate := temp_path / "tmp.crt").write_bytes(key.public_key)
        subprocess.run(
            ["sbverify", "--cert", certificate, output_path / "image.sig"],
            check=True,
            stdout=subprocess.DEVNULL,
        )
