path: root/v-i

#!/usr/bin/python3

import argparse
import glob
import logging
import os
import shutil
import subprocess
import sys
import tempfile
import yaml


verbose = False


def log(msg):
    if verbose:
        print("INSTALLER:", msg)
    logging.info(msg)


def run(argv, **kwargs):
    log(f"RUN: {argv} {kwargs}")
    return subprocess.run(argv, **kwargs)


def physical_volumes():
    log("list physical volumes")
    p = run(["pvdisplay", "-C", "--noheadings"], capture_output=True, check=True)
    lines = p.stdout.decode().splitlines()
    pvs = []
    for line in lines:
        words = line.split()
        pv = words[0]
        vg = words[1]
        pvs.append({"pv": pv, "vg": vg})
    return pvs


def logical_volumes():
    log("list logical volumes")
    p = run(["lvdisplay", "-C", "--noheadings"], capture_output=True, check=True)
    lines = p.stdout.decode().splitlines()
    lvs = []
    for line in lines:
        words = line.split()
        lv = words[0]
        vg = words[1]
        path = os.path.realpath(f"/dev/{vg}/{lv}")
        lvs.append({"lv": lv, "vg": vg, "path": path})
    return lvs


def volume_groups():
    log("list volume groups")
    p = run(["vgdisplay", "-C", "--noheadings"], capture_output=True, check=True)
    lines = p.stdout.decode().splitlines()
    return [line.split()[0] for line in lines if line.strip()]


def mount_points():
    log("find mount all points")
    mounts = {}
    for line in open("/proc/mounts").readlines():
        words = line.split()
        dev = os.path.realpath(words[0]) if words[0].startswith("/") else words[0]
        mounts[dev] = {
            "mount": words[1],
            "type": words[2],
        }
    return mounts


def find_mount_points(mounts, dev):
    log(f"find mount points using {dev}")
    res = []
    if dev in mounts:
        root = mounts[dev]["mount"]
        for m in mounts.values():
            if m["mount"].startswith(root):
                res.append(m)
    return res


def is_luks(path):
    log(f"is {path} a LUKS device?")
    p = run(["cryptsetup", "isLuks", path], check=False)
    return p.returncode == 0


def clean_up_disks(drives):
    log(f"clean system LVM2 and drives from everything: {drives}")

    mounts = mount_points()
    pvs = physical_volumes()
    lvs = logical_volumes()
    vgs = volume_groups()

    log(f"PVs: {pvs}")
    log(f"VGs: {vgs}")
    log(f"LVs: {lvs}")

    for lv in lvs:
        log(f"LV: {lv}")
        for m in find_mount_points(mounts, lv["path"]):
            log(f"unmount {m['mount']}")
            run(["umount", m["mount"]], capture_output=True)

    for vg in vgs:
        log(f"remove volume group {vg}")
        run(["vgremove", "--yes", vg], check=True, capture_output=True)

    for pv in pvs:
        log(f"remove physical volume {pv}")
        run(["pvremove", pv["pv"]], check=True, capture_output=True)
        if is_luks(pv["pv"]):
            run(["cryptsetup", "close", pv["pv"]], check=True, capture_output=True)

    for mapping in glob.glob("/dev/mapper/*"):
        if not mapping.endswith("/control"):
            log(f"open LUKS volume {mapping} (just in case it is one)")
            run(["cryptsetup", "close", mapping], check=False, capture_output=True)

    for drive in drives:
        log(f"blkdiscard {drive}")
        run(
            ["blkdiscard", "--force", drive],
            check=True,
            capture_output=True,
        )


def mklabel(device):
    return {
        # Create a GPT partition table. It works well for all modern
        # systems.
        "mklabel": "gpt",
        "device": device,
    }


def mkpart(device, tag, start, end):
    return {
        "mkpart": "primary",
        "device": device,
        "start": start,
        "end": end,
        "tag": tag,
    }


def mkfs(tag, fstype):
    return {
        "mkfs": fstype,
        "partition": tag,
    }


def cryptsetup(tag, password, name):
    return {
        "cryptsetup": "cryptsetup0",
        "password": password,
        "name": name,
    }


def vgcreate(tag, drives):
    return {
        "vgcreate": tag,
        "physical": drives,
    }


def lvcreate(vg, name, size):
    return {
        "lvcreate": vg,
        "name": name,
        "size": size,
    }


def mount(tag, dirname=None, mount_on=None):
    d = {"mount": tag, "zerofree": False}
    if dirname is not None:
        d["dirname"] = dirname
    if mount_on is not None:
        d["mount-on"] = mount_on
    return d


def unpack_rootfs(tag):
    return {
        "unpack-rootfs": tag,
    }


def cache_rootfs(tag):
    return {
        "cache-rootfs": tag,
        "unless": "rootfs_unpacked",
    }


def debootstrap(tag):
    return {
        "debootstrap": "bullseye",
        "mirror": "http://deb.debian.org/debian",
        "target": tag,
        "unless": "rootfs_unpacked",
        "require_empty_target": False,
    }


def apt(tag, packages, unless=True):
    d = {
        "apt": "install",
        "packages": packages,
        "tag": tag,
    }
    if unless:
        d["unless"] = "rootfs_unpacked"
    return d


def virtual_filesystems(tag):
    return {
        "virtual-filesystems": "root",
    }


def fstab(tag):
    return {
        "fstab": tag,
    }


def grub(device, root, efi):
    return {
        "grub": "uefi",
        "tag": root,
        "efi": efi,
        "quiet": True,
        "image-dev": device,
    }


def vmdb_spec(system, ansible_vars):
    device = "{{ image }}"
    steps = [
        mklabel(device),
        # Create a partition for UEFI to boot from.
        mkpart(device, "efi", "0%", "500M"),
        # Create a separate /boot partition, in case we want to have LUKS.
        mkpart(device, "boot", "500M", "1G"),
        # Format the EFI partition. This MUST be vfat.
        mkfs("efi", "vfat"),
        # Format /boot. This is conventionally ext2. This file system
        # gets very little I/O, but it MUST be supported by GRUB, so
        # ext2 seems like a nice, safe choice.
        mkfs("boot", "ext2"),
    ]

    # Set up pv0 for lvm2, either encrypted or cleartext.
    if system.luks:
        steps.extend(
            [
                mkpart(device, "cryptsetup0", "1G", "100%"),
                cryptsetup("cryptsetup0", system.luks, "pv0"),
            ]
        )
        for (i, drive) in enumerate(system.extra_drives):
            steps.append(cryptsetup(f"cryptsetuo{i+1}", system.luks, f"pv{i+1}"))
    else:
        steps.append(mkpart(device, "pv0", "1G", "100%"))
        for (i, drive) in enumerate(system.extra_drives):
            steps.extend(
                [
                    mklabel(drive),
                    mkpart(drive, f"pv{i+1}", "0%", "100%"),
                ]
            )

    # Create file systems and install Debian.
    steps.extend(
        [
            # Create an LVM2 volume group using pv0, which we create
            # earlier. At this point, if LUKS is used, pv0 is the unlocked,
            # open, cleartext block device.
            vgcreate(
                "vg0",
                [f"pv{i}" for i in range(len([system.drive] + system.extra_drives))],
            ),
            # Create a 20 gigabyte LV for the root file system. That's big
            # enough for a desktop system.
            lvcreate("vg0", "root", "20G"),
            # format the root file system. This gets a fair bit of use, so
            # ext4 seems like a safe choice. If you wanted another file
            # system, sorry.
            mkfs("root", "ext4"),
            # Mount the root file system.
            mount("root"),
            # Mount /boot on top of the root file system.
            mount("boot", dirname="/boot", mount_on="root"),
            # Mount /boot/efi.
            mount("efi", dirname="/boot/efi", mount_on="boot"),
        ]
    )

    # Add any additional LVs.
    for lv in system.extra_lvs:
        steps.extend(
            [
                lvcreate("vg0", lv["name"], lv["size"]),
                mkfs(lv["name"], "ext4"),
                mount(lv["name"], dirname=lv["mounted"], mount_on="root"),
            ]
        )

    steps.extend(
        [
            # If we have a cached version of the installed system, unpack
            # it now. Otherwise do nothing. Note that if you make any
            # changes to the steps marked "unless: rootfs_unpacked", you
            # have to remember to manually remove the cache file. v-i or
            # vmdb2 won't do that automatically for you.
            unpack_rootfs("root"),
            debootstrap("root"),
            apt(
                "root",
                [
                    "console-setup",
                    "dosfstools",
                    "linux-image-amd64",
                    "locales-all",
                    "lvm2",
                    "psmisc",
                    "python3",
                    "python3-apt",
                    "ssh",
                    "strace",
                ],
            ),
            # If we didn't unpack an existing cache archive, make one now.
            # Otherwise, skip this step.
            cache_rootfs("root"),
            # This MUST be after the debootstrap step.
            virtual_filesystems("root"),
            # Create /etc/fstab (and, if LUKS is used, /etc/crypttab).
            fstab("root"),
            # These MUST come after the fstab step so that they add the
            # crypttab in the initramfs. We install them regardless of
            # whether LUKS is used: they're harmless if LUKS isn't used.
            apt(
                "root",
                [
                    "cryptsetup",
                    "cryptsetup-initramfs",
                ],
                unless=False,
            ),
            # This also MUST come outside the rootfs caching, as it install
            # things outside the file systems, and those won't be in the
            # cache.
            grub(device, "root", "efi"),
        ]
    )

    # If playbooks have been specified, add ansible steps.
    for p in ["std.yml"] + system.extra_playbooks:
        if p:
            steps.append({"ansible": "root", "playbook": p, "extra_vars": ansible_vars})

    return {"steps": steps}


class SystemSpec:
    def __init__(self, filename):
        REQUIRED = "required"
        self._obj = {
            "hostname": REQUIRED,
            "drive": REQUIRED,
            "extra_drives": [],
            "extra_lvs": [],
            "extra_playbooks": [],
            "ansible_vars": {},
            "luks": "",
        }
        with open(filename) as f:
            obj = yaml.safe_load(f)

        # Check for unknown keys.
        for key in obj:
            if key not in self._obj:
                sys.exit(f"spec has unknown key: {key}")

        # Check for missing required keys.
        for key in self._obj:
            if self._obj[key] == REQUIRED and key not in obj:
                sys.exit(f"spec lacks required key {key}")

        # Check for types of values.
        for key in self._obj:
            if key in obj:
                e = type(self._obj[key])
                a = type(obj[key])
                if a != e:
                    sys.exit(f"spec key {key} has unexpected type {a}, wanted {e}")

        self._obj.update(obj)

        for key in self._obj:
            setattr(self, key, self._obj[key])
        del self._obj

    def __repr__(self):
        r = {key: getattr(self, key) for key in dir(self) if not key.startswith("_")}
        return repr(r)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--verbose", action="store_true")
    p.add_argument("--very-verbose", action="store_true")
    p.add_argument("--log", default="install.log")
    p.add_argument("--cache", default="cache.tar.gz")
    p.add_argument("spec")
    args = p.parse_args()

    logging.basicConfig(
        filename=args.log,
        level=logging.DEBUG,
        format="%(asctime)s %(levelname)s %(message)s",
    )

    global verbose
    verbose = args.verbose

    log("v-i starts")

    system = SystemSpec(args.spec)
    log(f"spec: {system!r}")

    clean_up_disks([system.drive] + system.extra_drives)

    ansible_vars = dict(system.ansible_vars)
    ansible_vars["hostname"] = system.hostname
    vmdb = vmdb_spec(system, ansible_vars)
    tmp = tempfile.mkdtemp()
    specfile = os.path.join(tmp, "spec.yaml")
    if args.very_verbose:
        yaml.dump(vmdb, stream=sys.stdout, indent=4)
    with open(specfile, "w") as f:
        yaml.dump(vmdb, stream=f, indent=4)

    log(f"run vmdb2 to install on {system.drive}")
    env = dict(os.environ)
    env["ANSIBLE_STDOUT_CALLBACK"] = "yaml"
    env["ANSIBLE_NOCOWS"] = "1"
    env["ANSIBLE_LOG_PATH"] = "ansible.log"

    argv = [
        "vmdb2",
        f"--rootfs-tarball={args.cache}",
        f"--log={args.log}",
        f"--image={system.drive}",
        specfile,
    ]
    if verbose:
        argv.append("--verbose")
    run(argv, check=True, capture_output=True)

    log("cleanup")
    shutil.rmtree(tmp)

    log("OK, done")
    print("OK, done")


main()