omarchy-iso/test/test_oem_state.py

"""Unit tests for OEM install support in the orchestrator.

Covers the InstallContext OEM plumbing (marker file, credential-less input,
throwaway LUKS passphrase injection) and the stage_oem_state /
create_factory_snapshot / configure_login phases against a temp target with
subprocess mocked out.
"""

import json
import os
import sys
import tempfile
import types
import unittest
from pathlib import Path
from subprocess import CompletedProcess
from unittest import mock

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "configs/airootfs/usr/share/omarchy-iso"))

sys.modules.setdefault(
    "orchestrator.archinstall_adapter", types.ModuleType("orchestrator.archinstall_adapter")
)

from orchestrator import phases_impl  # noqa: E402
from orchestrator.context import InstallContext  # noqa: E402


class ContextOemTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        self.dir = Path(self.tmp.name)
        self.state_dir = self.dir / "state"

        self.env = {
            "OMARCHY_INSTALL_CONFIG": str(self.dir / "user_configuration.json"),
            "OMARCHY_INSTALL_CREDS": str(self.dir / "user_credentials.json"),
            "OMARCHY_INSTALL_STATE_DIR": str(self.state_dir),
        }

    def write_config(self, config):
        (self.dir / "user_configuration.json").write_text(json.dumps(config))

    def write_creds(self, creds):
        (self.dir / "user_credentials.json").write_text(json.dumps(creds))

    def from_env(self, **extra_env):
        env = {**self.env, **extra_env}
        with mock.patch.dict(os.environ, env, clear=False):
            for key in ("OMARCHY_INSTALL_OEM_FILE",):
                if key not in env:
                    os.environ.pop(key, None)
            return InstallContext.from_env()

    def base_config(self, oem=None, disk_encryption=None):
        config = {
            "disk_config": {"config_type": "default_layout"},
            "omarchy_install": {"mode": "full_disk", "target_mount": "/mnt"},
        }
        if oem is not None:
            config["omarchy_install"]["oem"] = oem
        if disk_encryption is not None:
            config["disk_config"]["disk_encryption"] = disk_encryption
        return config

    def test_oem_flag_from_config(self):
        self.write_config(self.base_config(oem=True))
        ctx = self.from_env()
        self.assertTrue(ctx.oem)
        self.assertEqual(ctx.user_credentials, {"users": []})
        self.assertEqual(ctx.username, "")

    def test_oem_marker_file_arms_oem_without_credentials(self):
        self.write_config(self.base_config())
        marker = self.dir / "oem"
        marker.touch()
        ctx = self.from_env(OMARCHY_INSTALL_OEM_FILE=str(marker))
        self.assertTrue(ctx.oem)
        self.assertTrue(ctx.omarchy_install["oem"])

    def test_absent_marker_file_is_not_oem(self):
        self.write_config(self.base_config())
        self.write_creds({"users": [{"username": "jeff"}]})
        ctx = self.from_env(OMARCHY_INSTALL_OEM_FILE=str(self.dir / "missing-oem"))
        self.assertFalse(ctx.oem)
        self.assertEqual(ctx.username, "jeff")

    def test_missing_credentials_without_oem_raises(self):
        self.write_config(self.base_config())
        with self.assertRaisesRegex(RuntimeError, "credentials file missing"):
            self.from_env()

    def test_oem_generates_throwaway_luks_passphrase(self):
        self.write_config(self.base_config(oem=True, disk_encryption={
            "encryption_type": "luks", "partitions": ["x"],
        }))
        ctx = self.from_env()

        password = ctx.user_configuration["disk_config"]["disk_encryption"]["encryption_password"]
        self.assertTrue(password)
        self.assertEqual(ctx.user_credentials["encryption_password"], password)

        # archinstall reads both files from disk; the injected passphrase must
        # be present in each.
        arch_config = json.loads(ctx.arch_config_path.read_text())
        self.assertEqual(
            arch_config["disk_config"]["disk_encryption"]["encryption_password"], password
        )
        synthesized = json.loads(ctx.creds_path.read_text())
        self.assertEqual(synthesized["encryption_password"], password)

    def test_oem_reuses_rig_supplied_passphrase(self):
        self.write_config(self.base_config(oem=True, disk_encryption={
            "encryption_type": "luks", "partitions": ["x"],
        }))
        self.write_creds({"users": [], "encryption_password": "rig-secret"})
        ctx = self.from_env()
        self.assertEqual(
            ctx.user_configuration["disk_config"]["disk_encryption"]["encryption_password"],
            "rig-secret",
        )

    def test_oem_unencrypted_touches_nothing(self):
        self.write_config(self.base_config(oem=True))
        ctx = self.from_env()
        self.assertNotIn("encryption_password", ctx.user_credentials)


def make_ctx(target, **overrides):
    defaults = dict(
        target=target,
        oem=True,
        encrypt=False,
        username="",
        omarchy_install={"mode": "full_disk", "oem": True, "storage": {}},
        user_configuration={"disk_config": {}},
        user_credentials={"users": []},
        state_dir=target / "state",
        authorized_keys_path=None,
    )
    defaults.update(overrides)
    return types.SimpleNamespace(**defaults)


class StageOemStateTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        self.target = Path(self.tmp.name) / "mnt"
        self.target.mkdir()

        info_patch = mock.patch.object(phases_impl, "info")
        info_patch.start()
        self.addCleanup(info_patch.stop)

        # A fake bundled Node tarball on the "live ISO".
        self.packages = Path(self.tmp.name) / "opt-packages"
        self.packages.mkdir()
        (self.packages / "node-v24.0.0-linux-x64.tar.gz").write_bytes(b"node")
        node_patch = mock.patch.object(phases_impl, "NODE_PACKAGES_DIR", self.packages)
        node_patch.start()
        self.addCleanup(node_patch.stop)

    def install_runtime_oem_support(self):
        service = self.target / "usr/share/omarchy/install/oem/omarchy-oem-setup.service"
        service.parent.mkdir(parents=True, exist_ok=True)
        service.write_text("[Unit]\n")
        setup_bin = self.target / "usr/bin/omarchy-oem-setup"
        setup_bin.parent.mkdir(parents=True, exist_ok=True)
        setup_bin.write_text("#!/bin/bash\n")

    def oem_dir(self):
        return self.target / "var/lib/omarchy/oem"

    def test_non_oem_install_stages_only_the_node_tarball(self):
        ctx = make_ctx(self.target, oem=False)
        phases_impl.stage_oem_state(ctx)

        self.assertTrue((self.oem_dir() / "packages/node-v24.0.0-linux-x64.tar.gz").exists())
        self.assertFalse((self.oem_dir() / "pending").exists())
        self.assertFalse((self.target / "etc/systemd/system/omarchy-oem-setup.service").exists())

    def test_oem_against_runtime_without_support_fails_clearly(self):
        ctx = make_ctx(self.target)
        with self.assertRaisesRegex(RuntimeError, "does not ship"):
            phases_impl.stage_oem_state(ctx)

    def test_oem_without_node_tarball_fails(self):
        for tarball in self.packages.glob("*.tar.gz"):
            tarball.unlink()
        self.install_runtime_oem_support()
        ctx = make_ctx(self.target)
        with self.assertRaisesRegex(RuntimeError, "Node tarball"):
            phases_impl.stage_oem_state(ctx)

    def test_oem_arms_first_boot_setup(self):
        self.install_runtime_oem_support()
        ctx = make_ctx(self.target)
        phases_impl.stage_oem_state(ctx)

        self.assertTrue((self.oem_dir() / "pending").exists())
        self.assertTrue((self.target / "etc/systemd/system/omarchy-oem-setup.service").exists())
        link = self.target / "etc/systemd/system/multi-user.target.wants/omarchy-oem-setup.service"
        self.assertTrue(link.is_symlink())
        self.assertEqual(
            os.readlink(link), "/etc/systemd/system/omarchy-oem-setup.service"
        )
        # Unencrypted: no LUKS staging.
        self.assertFalse((self.oem_dir() / "luks-key").exists())

    def test_oem_encrypted_stages_luks_auto_unlock(self):
        self.install_runtime_oem_support()
        ctx = make_ctx(
            self.target,
            user_configuration={"disk_config": {"disk_encryption": {
                "encryption_type": "luks",
                "encryption_password": "throwaway-secret",
            }}},
        )
        phases_impl.stage_oem_state(ctx)

        # Byte-for-byte the slot passphrase — no trailing newline.
        self.assertEqual((self.oem_dir() / "luks-key").read_text(), "throwaway-secret")
        keyfile = self.target / "etc/omarchy/oem.key"
        self.assertEqual(keyfile.read_text(), "throwaway-secret")
        self.assertEqual(keyfile.stat().st_mode & 0o777, 0o600)

        cmdline = (self.target / "etc/limine-entry-tool.d/99-omarchy-oem-unlock.conf").read_text()
        self.assertIn("cryptkey=rootfs:/etc/omarchy/oem.key", cmdline)
        files = (self.target / "etc/mkinitcpio.conf.d/99-omarchy-oem-key.conf").read_text()
        self.assertIn("FILES+=(/etc/omarchy/oem.key)", files)

    def test_oem_pre_encrypted_without_passphrase_fails(self):
        self.install_runtime_oem_support()
        ctx = make_ctx(
            self.target,
            omarchy_install={"mode": "protected", "oem": True, "storage": {"luks_uuid": "abc"}},
        )
        with self.assertRaisesRegex(RuntimeError, "passphrase"):
            phases_impl.stage_oem_state(ctx)


class ConfigureLoginOemTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        self.target = Path(self.tmp.name)
        self.calls = []

        run_patch = mock.patch.object(
            phases_impl.subprocess, "run",
            side_effect=lambda cmd, **kw: (self.calls.append(cmd), CompletedProcess(cmd, 0))[1],
        )
        run_patch.start()
        self.addCleanup(run_patch.stop)

    def test_oem_login_leaves_user_state_to_first_boot(self):
        ctx = make_ctx(self.target, encrypt=True)
        phases_impl.configure_login(ctx)

        self.assertTrue((self.target / "etc/sddm.conf.d/99-omarchy-login.conf").exists())
        self.assertFalse((self.target / "etc/sddm.conf.d/autologin.conf").exists())
        self.assertFalse((self.target / "var/lib/sddm/state.conf").exists())
        self.assertTrue(any("sddm.service" in cmd for cmd in self.calls))

    def test_normal_encrypted_login_still_autologs_in(self):
        ctx = make_ctx(self.target, oem=False, encrypt=True, username="jeff")
        phases_impl.configure_login(ctx)

        autologin = (self.target / "etc/sddm.conf.d/autologin.conf").read_text()
        self.assertIn("User=jeff", autologin)
        state = (self.target / "var/lib/sddm/state.conf").read_text()
        self.assertIn("User=jeff", state)


class ConfigureSshAccessOemTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        self.target = Path(self.tmp.name) / "mnt"
        self.target.mkdir()

        info_patch = mock.patch.object(phases_impl, "info")
        info_patch.start()
        self.addCleanup(info_patch.stop)

        def fake_run(cmd, **kwargs):
            self.calls.append(cmd)
            if "ufw" in cmd:
                rules = self.target / "etc/ufw/user.rules"
                rules.parent.mkdir(parents=True, exist_ok=True)
                rules.write_text("-A ufw-user-input -p tcp --dport 22 -j ACCEPT\n")
            return CompletedProcess(cmd, 0)

        self.calls = []
        run_patch = mock.patch.object(phases_impl.subprocess, "run", side_effect=fake_run)
        run_patch.start()
        self.addCleanup(run_patch.stop)

    def test_oem_stages_keys_for_first_boot(self):
        keys = Path(self.tmp.name) / "authorized_keys"
        keys.write_text("ssh-ed25519 AAAA rig@host\n")
        ctx = make_ctx(self.target, authorized_keys_path=keys)

        phases_impl.configure_ssh_access(ctx)

        staged = self.target / "var/lib/omarchy/oem/authorized_keys"
        self.assertEqual(staged.read_text(), "ssh-ed25519 AAAA rig@host\n")
        self.assertEqual(staged.stat().st_mode & 0o777, 0o600)
        # No user yet — nothing under /home, no chown.
        self.assertFalse((self.target / "home").exists())
        self.assertFalse(any("chown" in cmd for cmd in self.calls))
        # The door is still opened for first boot.
        self.assertTrue(any("sshd.service" in cmd for cmd in self.calls))


class CreateFactorySnapshotTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        self.target = Path(self.tmp.name) / "mnt"
        self.target.mkdir()
        self.state_dir = Path(self.tmp.name) / "state"
        self.state_dir.mkdir()

        info_patch = mock.patch.object(phases_impl, "info")
        info_patch.start()
        self.addCleanup(info_patch.stop)

        self.findmnt = {
            "FSTYPE": "btrfs",
            "OPTIONS": "rw,noatime,compress=zstd:3,subvol=/@",
            "SOURCE": "/dev/mapper/omarchy_root[/@]",
        }
        self.calls = []

        def fake_run(cmd, **kwargs):
            self.calls.append(cmd)
            if cmd[0] == "findmnt":
                return CompletedProcess(cmd, 0, stdout=self.findmnt[cmd[2]] + "\n", stderr="")
            if cmd[0] == "mount":
                # Simulate the top level containing @.
                (Path(cmd[-1]) / "@").mkdir(parents=True, exist_ok=True)
            return CompletedProcess(cmd, 0, stdout="", stderr="")

        run_patch = mock.patch.object(phases_impl.subprocess, "run", side_effect=fake_run)
        run_patch.start()
        self.addCleanup(run_patch.stop)

    def ctx(self):
        return make_ctx(self.target, oem=False, state_dir=self.state_dir)

    def test_snapshots_root_as_factory(self):
        phases_impl.create_factory_snapshot(self.ctx())

        top = str(self.state_dir / "factory-top")
        self.assertIn(["mount", "-o", "subvolid=5", "/dev/mapper/omarchy_root", top], self.calls)
        self.assertIn(
            ["btrfs", "subvolume", "snapshot", "-r", f"{top}/@", f"{top}/@factory"],
            self.calls,
        )
        self.assertIn(["umount", top], self.calls)

    def test_non_btrfs_target_skips(self):
        self.findmnt["FSTYPE"] = "ext4"
        phases_impl.create_factory_snapshot(self.ctx())
        self.assertFalse(any(cmd[0] == "mount" for cmd in self.calls))

    def test_non_subvol_root_skips(self):
        self.findmnt["OPTIONS"] = "rw,noatime"
        phases_impl.create_factory_snapshot(self.ctx())
        self.assertFalse(any(cmd[0] == "mount" for cmd in self.calls))


if __name__ == "__main__":
    unittest.main()