from __future__ import annotations

import hashlib
import tempfile
import unittest
from pathlib import Path

from uvm.config import Settings
from uvm.errors import UvmError
from uvm.images import (
    _FIRECRACKER_DEMO_PUBLIC_KEY,
    _enable_ssh_password_authentication,
    _set_shadow_password,
    _without_firecracker_demo_key,
    ImageStore,
)
from uvm.integrity import write_manifest
from uvm.system import CommandResult


class PasswordHashRunner:
    def __init__(self) -> None:
        self.command: tuple[str, ...] | None = None
        self.options: dict[str, object] = {}

    def run(self, command, **options) -> CommandResult:
        self.command = tuple(str(part) for part in command)
        self.options = options
        return CommandResult(self.command, 0, stdout="$6$salt$password-hash\n")


class ImageStoreTests(unittest.TestCase):
    def test_uses_the_persisted_install_manifest_when_strict_mode_is_enabled(self) -> None:
        with tempfile.TemporaryDirectory() as temporary_directory:
            settings = Settings(base=Path(temporary_directory), allow_unverified_downloads=False)
            settings.images_dir.mkdir(parents=True)
            settings.kernel_image.write_bytes(b"kernel")
            settings.rootfs_image.write_bytes(b"rootfs")
            write_manifest(
                settings.integrity_manifest_path,
                {
                    "kernel": hashlib.sha256(b"kernel").hexdigest(),
                    "rootfs": hashlib.sha256(b"rootfs").hexdigest(),
                    "firecracker": "0" * 64,
                    "jailer": "1" * 64,
                },
                verified=True,
            )

            assets = ImageStore(settings).installed_assets()

        self.assertEqual(assets.kernel.name, "vmlinux")
        self.assertEqual(assets.rootfs.name, "ubuntu.ext4")

    def test_hashes_password_through_stdin_without_putting_it_in_argv(self) -> None:
        runner = PasswordHashRunner()
        store = ImageStore(Settings(), runner=runner)  # type: ignore[arg-type]

        password_hash = store._password_hash("secret-value")

        self.assertEqual(password_hash, "$6$salt$password-hash")
        self.assertEqual(runner.command, ("openssl", "passwd", "-6", "-stdin"))
        self.assertNotIn("secret-value", runner.command)
        self.assertEqual(runner.options["input_text"], "secret-value\n")
        self.assertTrue(runner.options["sensitive"])

    def test_updates_only_the_requested_shadow_entry(self) -> None:
        original = "root:*:1:0:99999:7:::\nservice:!:1:0:99999:7:::\n"

        updated = _set_shadow_password(original, "root", "$6$salt$hash")

        self.assertEqual(
            updated,
            "root:$6$salt$hash:1:0:99999:7:::\nservice:!:1:0:99999:7:::\n",
        )

        with self.assertRaises(UvmError):
            _set_shadow_password(original, "missing", "$6$salt$hash")

    def test_enables_root_password_login_before_existing_sshd_settings(self) -> None:
        updated = _enable_ssh_password_authentication(
            "PasswordAuthentication no\nPermitRootLogin prohibit-password\n",
            "root",
        )

        self.assertTrue(
            updated.startswith(
                "# Managed by uvm\nPasswordAuthentication yes\nPermitRootLogin yes\n"
            )
        )

    def test_removes_only_the_public_firecracker_demo_key(self) -> None:
        own_key = "ssh-ed25519 AAAA-own-key developer@example"

        updated = _without_firecracker_demo_key(
            f"{_FIRECRACKER_DEMO_PUBLIC_KEY} demo\n{own_key}\n"
        )

        self.assertEqual(updated, f"{own_key}\n")