nest/tests/test_images.py

from __future__ import annotations

import hashlib
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from uvm.config import Settings
from uvm.errors import UvmError
from uvm.images import (
    _FIRECRACKER_DEMO_PUBLIC_KEY,
    _GuestFile,
    _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.calls: list[tuple[str, ...]] = []
        self.options: dict[str, object] = {}

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


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(
            "HostKey /insecure/shared-key\n"
            "PasswordAuthentication no\n"
            "PermitRootLogin prohibit-password\n",
            "root",
        )

        self.assertTrue(
            updated.startswith(
                "# Managed by uvm\n"
                "HostKey /etc/ssh/ssh_host_rsa_key\n"
                "HostKey /etc/ssh/ssh_host_ecdsa_key\n"
                "HostKey /etc/ssh/ssh_host_ed25519_key\n"
                "PasswordAuthentication no\n"
                "KbdInteractiveAuthentication no\n"
                "ChallengeResponseAuthentication no\n"
                "PermitRootLogin yes\n"
                "Match User root\n"
                "    PasswordAuthentication yes\n"
                "Match all\n"
            )
        )
        self.assertNotIn("/insecure/shared-key", updated)
        self.assertNotIn("PermitRootLogin prohibit-password", updated)

    def test_non_root_credentials_preserve_the_existing_root_login_policy(self) -> None:
        updated = _enable_ssh_password_authentication(
            "PermitRootLogin no\nChallengeResponseAuthentication yes\n",
            "service",
        )

        self.assertIn("PermitRootLogin no", updated)
        self.assertIn("Match User service\n    PasswordAuthentication yes", updated)
        self.assertNotIn("ChallengeResponseAuthentication yes", updated)

    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")

    def test_generates_three_unique_guest_host_key_types(self) -> None:
        runner = PasswordHashRunner()
        store = ImageStore(Settings(), runner=runner)  # type: ignore[arg-type]
        with (
            tempfile.TemporaryDirectory() as temporary_directory,
            patch.object(store, "_read_guest_file", return_value=None),
            patch.object(store, "_write_guest_file") as write_guest_file,
        ):
            store._replace_guest_host_keys(
                Path("/tmp/rootfs.ext4"),
                Path(temporary_directory),
            )

        keygen_calls = [call for call in runner.calls if call[0] == "ssh-keygen"]
        self.assertEqual([call[3] for call in keygen_calls], ["rsa", "ecdsa", "ed25519"])
        self.assertEqual(write_guest_file.call_count, 6)

    def test_optional_guest_file_inspection_fails_closed_on_debugfs_errors(self) -> None:
        class FailedRunner:
            @staticmethod
            def run(_command, **_options) -> CommandResult:
                return CommandResult(
                    args=("debugfs",),
                    returncode=0,
                    stderr="filesystem checksum failure",
                )

        store = ImageStore(Settings(), runner=FailedRunner())  # type: ignore[arg-type]
        with tempfile.TemporaryDirectory() as temporary_directory:
            with self.assertRaises(UvmError):
                store._read_guest_file(
                    Path("/tmp/rootfs.ext4"),
                    "/root/.ssh/authorized_keys",
                    Path(temporary_directory) / "authorized_keys",
                    required=False,
                )

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

        for guest_file in (
            _GuestFile(Path("/tmp/shadow"), 0o640, 0, 42, links=2),
            _GuestFile(
                Path("/tmp/shadow"),
                0o640,
                0,
                42,
                has_extended_attributes=True,
            ),
        ):
            with self.subTest(guest_file=guest_file):
                with self.assertRaises(UvmError):
                    store._write_guest_file(
                        Path("/tmp/rootfs.ext4"),
                        "/etc/shadow",
                        guest_file,
                    )

        self.assertEqual(runner.calls, [])