clan-core/pkgs/clan-cli/clan_cli/machines/install.py

import argparse
import importlib
import json
import logging
import os
from dataclasses import dataclass
from pathlib import Path
from tempfile import TemporaryDirectory

from ..clan_uri import FlakeId
from ..cmd import Log, run
from ..completions import add_dynamic_completer, complete_machines
from ..facts.generate import generate_facts
from ..machines.machines import Machine
from ..nix import nix_shell
from ..ssh.cli import is_ipv6, is_reachable, qrcode_scan

log = logging.getLogger(__name__)


class ClanError(Exception):
    pass


def install_nixos(
    machine: Machine,
    kexec: str | None = None,
    debug: bool = False,
    password: str | None = None,
    no_reboot: bool = False,
    extra_args: list[str] = [],
) -> None:
    secret_facts_module = importlib.import_module(machine.secret_facts_module)
    log.info(f"installing {machine.name}")
    secret_facts_store = secret_facts_module.SecretStore(machine=machine)

    h = machine.target_host
    target_host = f"{h.user or 'root'}@{h.host}"
    log.info(f"target host: {target_host}")

    generate_facts([machine], None, False)

    with TemporaryDirectory() as tmpdir_:
        tmpdir = Path(tmpdir_)
        upload_dir_ = machine.secrets_upload_directory

        if upload_dir_.startswith("/"):
            upload_dir_ = upload_dir_[1:]
        upload_dir = tmpdir / upload_dir_
        upload_dir.mkdir(parents=True)
        secret_facts_store.upload(upload_dir)

        if password:
            os.environ["SSHPASS"] = password

        cmd = [
            "nixos-anywhere",
            "--flake",
            f"{machine.flake}#{machine.name}",
            "--extra-files",
            str(tmpdir),
            *extra_args,
        ]

        if no_reboot:
            cmd.append("--no-reboot")

        if password:
            cmd += [
                "--env-password",
                "--ssh-option",
                "IdentitiesOnly=yes",
            ]

        if machine.target_host.port:
            cmd += ["--ssh-port", str(machine.target_host.port)]
        if kexec:
            cmd += ["--kexec", kexec]
        if debug:
            cmd.append("--debug")
        cmd.append(target_host)

        run(
            nix_shell(
                ["nixpkgs#nixos-anywhere"],
                cmd,
            ),
            log=Log.BOTH,
        )


@dataclass
class InstallOptions:
    flake: FlakeId
    machine: str
    target_host: str
    kexec: str | None
    confirm: bool
    debug: bool
    no_reboot: bool
    json_ssh_deploy: dict[str, str] | None
    nix_options: list[str]


def install_command(args: argparse.Namespace) -> None:
    json_ssh_deploy = None
    if args.json:
        json_file = Path(args.json)
        if json_file.is_file():
            json_ssh_deploy = json.loads(json_file.read_text())
        else:
            json_ssh_deploy = json.loads(args.json)
    elif args.png:
        json_ssh_deploy = json.loads(qrcode_scan(args.png))

    if not json_ssh_deploy and not args.target_host:
        raise ClanError("No target host provided, please provide a target host.")

    if json_ssh_deploy:
        target_host = f"root@{find_reachable_host_from_deploy_json(json_ssh_deploy)}"
        password = json_ssh_deploy["pass"]
    else:
        target_host = args.target_host
        password = None

    opts = InstallOptions(
        flake=FlakeId(args.flake),
        machine=args.machine,
        target_host=target_host,
        kexec=args.kexec,
        confirm=not args.yes,
        debug=args.debug,
        no_reboot=args.no_reboot,
        json_ssh_deploy=json_ssh_deploy,
        nix_options=args.option,
    )
    machine = Machine(opts.machine, flake=opts.flake)
    machine.target_host_address = opts.target_host

    if opts.confirm:
        ask = input(f"Install {machine.name} to {opts.target_host}? [y/N] ")
        if ask != "y":
            return

    install_nixos(
        machine,
        kexec=opts.kexec,
        debug=opts.debug,
        password=password,
        no_reboot=opts.no_reboot,
        extra_args=opts.nix_options,
    )


def find_reachable_host_from_deploy_json(deploy_json: dict[str, str]) -> str:
    host = None
    for addr in deploy_json["addrs"]:
        if is_reachable(addr):
            if is_ipv6(addr):
                host = f"[{addr}]"
            else:
                host = addr
            break
    if not host:
        raise ClanError(
            f"""
            Could not reach any of the host addresses provided in the json string.
            Please doublecheck if they are reachable from your machine.
            Try `ping [ADDR]` with one of the addresses: {deploy_json['addrs']}
            """
        )
    return host


def register_install_parser(parser: argparse.ArgumentParser) -> None:
    parser.add_argument(
        "--kexec",
        type=str,
        help="use another kexec tarball to bootstrap NixOS",
    )
    parser.add_argument(
        "--no-reboot",
        action="store_true",
        help="do not reboot after installation",
        default=False,
    )
    parser.add_argument(
        "--yes",
        action="store_true",
        help="do not ask for confirmation",
        default=False,
    )

    machines_parser = parser.add_argument(
        "machine",
        type=str,
        help="machine to install",
    )
    add_dynamic_completer(machines_parser, complete_machines)

    parser.add_argument(
        "target_host",
        type=str,
        nargs="?",
        help="ssh address to install to in the form of user@host:2222",
    )
    group = parser.add_mutually_exclusive_group(required=False)
    group.add_argument(
        "-j",
        "--json",
        help="specify the json file for ssh data (generated by starting the clan installer)",
    )
    group.add_argument(
        "-P",
        "--png",
        help="specify the json file for ssh data as the qrcode image (generated by starting the clan installer)",
    )
    parser.set_defaults(func=install_command)