clan-core/pkgs/clan-cli/clan_cli/vms/create.py

import argparse
import asyncio
import json
import os
import re
import shlex
import sys
import tempfile
from pathlib import Path
from typing import Iterator
from uuid import UUID

from ..dirs import clan_flakes_dir, specific_flake_dir
from ..nix import nix_build, nix_config, nix_eval, nix_shell
from ..task_manager import BaseTask, Command, create_task
from ..types import validate_path
from .inspect import VmConfig, inspect_vm


def is_flake_url(s: str) -> bool:
    if re.match(r"^http.?://[a-zA-Z0-9.-]+/[a-zA-Z0-9.-]+", s) is not None:
        return True
    return False


class BuildVmTask(BaseTask):
    def __init__(self, uuid: UUID, vm: VmConfig, nix_options: list[str] = []) -> None:
        super().__init__(uuid, num_cmds=7)
        self.vm = vm
        self.nix_options = nix_options

    def get_vm_create_info(self, cmds: Iterator[Command]) -> dict:
        config = nix_config()
        system = config["system"]

        clan_dir = self.vm.flake_url
        machine = self.vm.flake_attr
        cmd = next(cmds)
        cmd.run(
            nix_build(
                [
                    f'{clan_dir}#clanInternals.machines."{system}"."{machine}".config.system.clan.vm.create'
                ]
                + self.nix_options
            ),
            name="buildvm",
        )
        vm_json = "".join(cmd.stdout).strip()
        self.log.debug(f"VM JSON path: {vm_json}")
        with open(vm_json) as f:
            return json.load(f)

    def get_clan_name(self, cmds: Iterator[Command]) -> str:
        clan_dir = self.vm.flake_url
        cmd = next(cmds)
        cmd.run(
            nix_eval([f"{clan_dir}#clanInternals.clanName"]) + self.nix_options,
            name="clanname",
        )
        clan_name = cmd.stdout[0].strip().strip('"')
        return clan_name

    def run(self) -> None:
        cmds = self.commands()

        machine = self.vm.flake_attr
        self.log.debug(f"Creating VM for {machine}")

        # TODO: We should get this from the vm argument
        vm_config = self.get_vm_create_info(cmds)
        clan_name = self.get_clan_name(cmds)

        self.log.debug(f"Building VM for clan name: {clan_name}")

        flake_dir = clan_flakes_dir() / clan_name
        validate_path(clan_flakes_dir(), flake_dir)
        flake_dir.mkdir(exist_ok=True)

        with tempfile.TemporaryDirectory() as tmpdir_:
            tmpdir = Path(tmpdir_)
            xchg_dir = tmpdir / "xchg"
            xchg_dir.mkdir(exist_ok=True)
            secrets_dir = tmpdir / "secrets"
            secrets_dir.mkdir(exist_ok=True)
            disk_img = tmpdir / "disk.img"

            env = os.environ.copy()
            env["CLAN_DIR"] = str(self.vm.flake_url)

            env["PYTHONPATH"] = str(
                ":".join(sys.path)
            )  # TODO do this in the clanCore module
            env["SECRETS_DIR"] = str(secrets_dir)

            # Only generate secrets for local clans
            if not is_flake_url(str(self.vm.flake_url)):
                cmd = next(cmds)
                if Path(self.vm.flake_url).is_dir():
                    cmd.run(
                        [vm_config["generateSecrets"], clan_name],
                        env=env,
                        name="generateSecrets",
                    )
                else:
                    self.log.warning("won't generate secrets for non local clan")

            cmd = next(cmds)
            cmd.run(
                [vm_config["uploadSecrets"]],
                env=env,
                name="uploadSecrets",
            )

            cmd = next(cmds)
            cmd.run(
                nix_shell(
                    ["qemu"],
                    [
                        "qemu-img",
                        "create",
                        "-f",
                        "raw",
                        str(disk_img),
                        "1024M",
                    ],
                ),
                name="createDisk",
            )

            cmd = next(cmds)
            cmd.run(
                nix_shell(
                    ["e2fsprogs"],
                    [
                        "mkfs.ext4",
                        "-L",
                        "nixos",
                        str(disk_img),
                    ],
                ),
                name="formatDisk",
            )

            cmd = next(cmds)
            cmdline = [
                (Path(vm_config["toplevel"]) / "kernel-params").read_text(),
                f'init={vm_config["toplevel"]}/init',
                f'regInfo={vm_config["regInfo"]}/registration',
                "console=ttyS0,115200n8",
                "console=tty0",
            ]
            qemu_command = [
                "qemu-kvm",
                "-name",
                machine,
                "-m",
                f'{vm_config["memorySize"]}M',
                "-smp",
                str(vm_config["cores"]),
                "-device",
                "virtio-rng-pci",
                "-net",
                "nic,netdev=user.0,model=virtio",
                "-netdev",
                "user,id=user.0",
                "-virtfs",
                "local,path=/nix/store,security_model=none,mount_tag=nix-store",
                "-virtfs",
                f"local,path={xchg_dir},security_model=none,mount_tag=shared",
                "-virtfs",
                f"local,path={xchg_dir},security_model=none,mount_tag=xchg",
                "-virtfs",
                f"local,path={secrets_dir},security_model=none,mount_tag=secrets",
                "-drive",
                f"cache=writeback,file={disk_img},format=raw,id=drive1,if=none,index=1,werror=report",
                "-device",
                "virtio-blk-pci,bootindex=1,drive=drive1,serial=root",
                "-device",
                "virtio-keyboard",
                "-vga",
                "virtio",
                "-usb",
                "-device",
                "usb-tablet,bus=usb-bus.0",
                "-kernel",
                f'{vm_config["toplevel"]}/kernel',
                "-initrd",
                vm_config["initrd"],
                "-append",
                " ".join(cmdline),
            ]
            if not self.vm.graphics:
                qemu_command.append("-nographic")
            print("$ " + shlex.join(qemu_command))
            cmd.run(nix_shell(["qemu"], qemu_command), name="qemu")


def create_vm(vm: VmConfig, nix_options: list[str] = []) -> BuildVmTask:
    return create_task(BuildVmTask, vm, nix_options)


def create_command(args: argparse.Namespace) -> None:
    flake_url = args.flake
    if not is_flake_url(str(args.flake)):
        flake_url = specific_flake_dir(args.flake)
    vm = asyncio.run(inspect_vm(flake_url=flake_url, flake_attr=args.machine))

    task = create_vm(vm, args.option)
    for line in task.log_lines():
        print(line, end="")


def register_create_parser(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("machine", type=str, help="machine in the flake to create")
    parser.set_defaults(func=create_command)