# 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 signing models."""

import io
import os
import random
import re
import subprocess
from pathlib import Path
from typing import Any
from unittest import mock

import django.test

from nacl.public import PrivateKey

import debusine.test
from debusine.signing.models import (
    AuditLog,
    Key,
    ProtectedKey,
    ProtectedKeyError,
    ProtectedKeyNaCl,
    ProtectedKeyStorage,
)


class ProtectedKeyNaClTests(django.test.TestCase):
    """Unit tests for `ProtectedKeyNaCl`."""

    def test_encrypt_error(self) -> None:
        """Encryption errors raise ProtectedKeyError."""
        # This doesn't fail in normal use unless libsodium is broken.  Force
        # it to fail using a type error.
        public_key = PrivateKey.generate().public_key
        with self.assertRaisesRegex(
            ProtectedKeyError, "input message must be bytes"
        ):
            ProtectedKeyNaCl.encrypt(
                public_key,
                "should be bytes",  # type: ignore[arg-type]
            )

    def test_decrypt_wrong_key(self) -> None:
        """Decrypting with the wrong key raises ProtectedKeyError."""
        private_keys = [PrivateKey.generate() for _ in range(2)]
        protected = ProtectedKeyNaCl.encrypt(
            private_keys[0].public_key, b"data"
        )
        with self.assertRaisesRegex(
            ProtectedKeyError,
            "Key not encrypted using any of the given private keys",
        ):
            protected.decrypt([private_keys[1]])

    def test_decrypt_error(self) -> None:
        """Other decryption errors raise ProtectedKeyError."""
        private_key = PrivateKey.generate()
        protected = ProtectedKeyNaCl(
            public_key=bytes(private_key.public_key), encrypted=b"nonsense" * 8
        )
        with self.assertRaisesRegex(
            ProtectedKeyError, "An error occurred trying to decrypt the message"
        ):
            protected.decrypt([private_key])

    def test_encrypt_decrypt(self) -> None:
        """Encrypting and decrypting some data works."""
        private_key = PrivateKey.generate()
        protected = ProtectedKeyNaCl.encrypt(private_key.public_key, b"data")
        self.assertEqual(protected.decrypt([private_key]), b"data")

    def test_encrypt_decrypt_rotated(self) -> None:
        """The key-encryption key can be rotated."""
        private_keys = [PrivateKey.generate() for _ in range(2)]
        protected = ProtectedKeyNaCl.encrypt(
            private_keys[1].public_key, b"data"
        )
        self.assertEqual(protected.decrypt(private_keys), b"data")


def _make_fake_fingerprint(length=64):
    """Generate a random string in the style of a fingerprint."""
    return "".join(random.choices("0123456789ABCDEF", k=length))


class KeyManagerTests(django.test.TestCase):
    """Unit tests for `KeyManager`."""

    def test_generate_uefi(self) -> None:
        """Generate a UEFI key."""
        # Generating keys is slow, so we just test that the correct
        # subprocesses are executed.
        private_key = random.randbytes(64)
        public_key = random.randbytes(64)
        fingerprint = _make_fake_fingerprint()
        storage_private_key = PrivateKey.generate()
        log_file = io.BytesIO()

        def run(
            args: list[str | os.PathLike[str]], **kwargs: Any
        ) -> subprocess.CompletedProcess[Any]:
            match args[1]:
                case "req":
                    Path(args[args.index("-keyout") + 1]).write_bytes(
                        private_key
                    )
                    Path(args[args.index("-out") + 1]).write_bytes(public_key)
                    return subprocess.CompletedProcess(args=args, returncode=0)
                case "x509":
                    fp_colons = ":".join(re.findall("..", fingerprint))
                    return subprocess.CompletedProcess(
                        args=args,
                        returncode=0,
                        stdout=f"sha256 Fingerprint={fp_colons}\n".encode(),
                    )
                case _:  # pragma: no cover
                    return subprocess.CompletedProcess(args=args, returncode=1)

        with (
            mock.patch("subprocess.run", side_effect=run) as mock_run,
            self.settings(DEBUSINE_SIGNING_PRIVATE_KEYS=[storage_private_key]),
        ):
            key = Key.objects.generate(
                Key.Purpose.UEFI, "debian/bookworm test key", 123, log_file
            )

        self.assertEqual(mock_run.call_count, 2)
        mock_run.assert_has_calls(
            [
                mock.call(
                    [
                        "openssl",
                        "req",
                        "-new",
                        "-newkey",
                        "rsa:2048",
                        "-x509",
                        "-subj",
                        r"/CN=debian\/bookworm test key/",
                        "-days",
                        "5475",
                        "-noenc",
                        "-sha256",
                        "-keyout",
                        mock.ANY,
                        "-out",
                        mock.ANY,
                    ],
                    check=True,
                    stdout=log_file,
                    stderr=log_file,
                ),
                mock.call(
                    [
                        "openssl",
                        "x509",
                        "-inform",
                        "PEM",
                        "-noout",
                        "-fingerprint",
                        "-sha256",
                    ],
                    check=True,
                    input=mock.ANY,
                    stdout=subprocess.PIPE,
                    stderr=log_file,
                ),
            ]
        )
        self.assertEqual(key.purpose, Key.Purpose.UEFI)
        self.assertEqual(key.fingerprint, fingerprint)
        self.assertEqual(
            key.stored_private_key.storage, ProtectedKeyStorage.NACL
        )
        self.assertEqual(
            key.stored_private_key.data.decrypt([storage_private_key]),
            private_key,
        )
        self.assertEqual(key.public_key, public_key)
        audit_log = AuditLog.objects.latest("created_at")
        self.assertEqual(audit_log.purpose, Key.Purpose.UEFI)
        self.assertEqual(audit_log.fingerprint, fingerprint)
        self.assertEqual(audit_log.event, AuditLog.Event.GENERATE)
        self.assertEqual(
            audit_log.data, {"description": "debian/bookworm test key"}
        )
        self.assertEqual(audit_log.created_by_work_request_id, 123)


class KeyTests(debusine.test.TestCase, django.test.TestCase):
    """Unit tests for `Key`."""

    def test_sign_uefi(self) -> None:
        """Sign data using a UEFI key."""
        # Signing data is slow, so we just test that the correct
        # subprocesses are executed.
        private_key = random.randbytes(64)
        public_key = random.randbytes(64)
        fingerprint = _make_fake_fingerprint()
        storage_private_key = PrivateKey.generate()
        key = Key.objects.create(
            purpose=Key.Purpose.UEFI,
            fingerprint=fingerprint,
            private_key=ProtectedKey(
                storage=ProtectedKeyStorage.NACL,
                data=ProtectedKeyNaCl.encrypt(
                    storage_private_key.public_key, private_key
                ),
            ).dict(),
            public_key=public_key,
        )
        temp_path = self.create_temporary_directory()
        (data_path := temp_path / "data").write_bytes(b"data to sign")
        signature_path = temp_path / "signed"
        signature_bytes = random.randbytes(64)
        log_file = io.BytesIO()
        key_bytes: bytes | None = None
        certificate_bytes: bytes | None = None

        def run(
            args: list[str | os.PathLike[str]], **kwargs: Any
        ) -> subprocess.CompletedProcess[Any]:
            nonlocal key_bytes, certificate_bytes
            key_bytes = Path(args[args.index("--key") + 1]).read_bytes()
            certificate_bytes = Path(
                args[args.index("--cert") + 1]
            ).read_bytes()
            Path(args[args.index("--output") + 1]).write_bytes(signature_bytes)
            return subprocess.CompletedProcess(args=args, returncode=0)

        with (
            mock.patch("subprocess.run", side_effect=run) as mock_run,
            self.settings(DEBUSINE_SIGNING_PRIVATE_KEYS=[storage_private_key]),
        ):
            key.sign(data_path, signature_path, 123, log_file)

        mock_run.assert_called_once_with(
            [
                "sbsign",
                "--key",
                mock.ANY,
                "--cert",
                mock.ANY,
                "--output",
                signature_path,
                data_path,
            ],
            check=True,
            stdout=log_file,
            stderr=log_file,
        )
        self.assertEqual(key_bytes, private_key)
        self.assertEqual(certificate_bytes, public_key)
        self.assertEqual(signature_path.read_bytes(), signature_bytes)
