tasks.py

# ruff: noqa: ARG001, S603, S607
import json
import os
import subprocess
from pathlib import Path
from tempfile import TemporaryDirectory

from deploykit import DeployGroup, DeployHost, parse_hosts
from invoke import Context, task

ROOT = Path(__file__).parent.resolve()
os.chdir(ROOT)


@task
def deploy(c: Context, hosts: str) -> None:
    """
    Use inv deploy --hosts build01,darwin01
    """
    g = DeployGroup(get_hosts(hosts))

    def deploy(h: DeployHost) -> None:
        if "darwin" in h.host:
            command = "sudo -H darwin-rebuild"
            target = f"{h.user}@{h.host}"
        else:
            command = "sudo nixos-rebuild"
            target = f"{h.host}"

        res = subprocess.run(
            ["nix", "flake", "metadata", "--json"],
            text=True,
            check=True,
            stdout=subprocess.PIPE,
        )
        data = json.loads(res.stdout)
        path = data["path"]

        send = (
            "nix flake archive"
            if any(
                (
                    n.get("locked", {}).get("type") == "path"
                    or n.get("locked", {}).get("url", "").startswith("file:")
                )
                for n in data["locks"]["nodes"].values()
            )
            else f"nix copy {path}"
        )

        h.run_local(f"{send} --to ssh://{target}")

        hostname = h.host.replace(".nix-community.org", "")
        h.run(
            f"{command} switch --option accept-flake-config true --flake {path}#{hostname}"
        )

    g.run_function(deploy)


@task
def update_sops_files(c: Context) -> None:
    """
    Update all sops yaml files according to sops.nix rules
    """
    c.run(f"nix eval --json -f {ROOT}/sops.nix | yq e -P - > {ROOT}/.sops.yaml")
    c.run(
        "shopt -s globstar && sops updatekeys --yes **/secrets.yaml modules/secrets/*"
    )


@task
def print_keys(c: Context, flake_attr: str) -> None:
    """
    Decrypt host private key, print ssh public key. Use inv print-keys --flake-attr build01
    """
    with TemporaryDirectory() as tmpdir:
        decrypt_host_key(flake_attr, tmpdir)
        key = f"{tmpdir}/var/lib/ssh_secrets/ssh_host_ed25519_key"
        pubkey = subprocess.run(
            ["ssh-keygen", "-y", "-f", f"{key}"],
            stdout=subprocess.PIPE,
            text=True,
            check=True,
        )
        print("###### Public keys ######")
        print(pubkey.stdout)


@task
def docs(c: Context) -> None:
    """
    Serve docs (mkdoc serve)
    """
    c.run("nix develop .#mkdocs -c mkdocs serve")


@task
def docs_linkcheck(c: Context) -> None:
    """
    Run docs online linkchecker
    """
    c.run("nix run .#docs-linkcheck.online")


def get_hosts(hosts: str) -> list[DeployHost]:
    g = parse_hosts(
        hosts=hosts,
        domain_suffix=".nix-community.org",
    )
    for i in g.hosts:
        if i.host.startswith("darwin"):
            i.user = "customer"
    return g.hosts


def decrypt_host_key(flake_attr: str, tmpdir: str) -> None:
    def opener(path: str, flags: int) -> int:
        return os.open(path, flags, 0o400)

    t = Path(tmpdir)
    t.mkdir(parents=True, exist_ok=True)
    t.chmod(0o755)

    def decrypt(path: str, secret: str) -> None:
        file = t / path
        file.parent.mkdir(parents=True, exist_ok=True)
        with open(str(file), "w", opener=opener) as fh:
            subprocess.run(
                [
                    "sops",
                    "--extract",
                    secret,
                    "--decrypt",
                    f"{ROOT}/secrets.yaml",
                ],
                check=True,
                stdout=fh,
            )

    decrypt(
        "var/lib/ssh_secrets/ssh_host_ed25519_key",
        f'["ssh_host_ed25519_key"]["{flake_attr}"]',
    )
    decrypt(
        "var/lib/ssh_secrets/initrd_host_ed25519_key", '["initrd_host_ed25519_key"]'
    )


@task
def install(c: Context, flake_attr: str, hostname: str) -> None:
    """
    Decrypt host private key, install with nixos-anywhere. Use inv install --flake-attr build01 --hostname build01.nix-community.org
    """
    ask = input(f"Install {hostname} with {flake_attr}? [y/N] ")
    if ask != "y":
        return
    with TemporaryDirectory() as tmpdir:
        decrypt_host_key(flake_attr, tmpdir)
        flags = "--build-on remote --debug --option accept-flake-config true"
        c.run(
            f"nix run --inputs-from . nixpkgs#nixos-anywhere -- {hostname} --extra-files {tmpdir} --flake .#{flake_attr} {flags}",
            echo=True,
        )


@task
def cleanup_gcroots(c: Context, hosts: str) -> None:
    g = DeployGroup(get_hosts(hosts))
    g.run("sudo find /nix/var/nix/gcroots/auto -type s -delete")