abroot-adapter

#!/bin/python3
import sys
import os
import uuid
import subprocess
import logging
import argparse
import json


def find_partition_used_space(partition: str) -> int:
    _part = "/dev/" + partition.split("/")[-1]
    _uuid = str(uuid.uuid4())
    subprocess.check_call(["mkdir", "-p", "/tmp/{}".format(_uuid)])
    subprocess.check_call(["sudo", "mount", _part, "/tmp/{}".format(_uuid)])
    _res = (
        subprocess.check_output(["df", "/tmp/{}".format(_uuid)])
        .decode("utf-8")
        .splitlines()[1]
        .split()
    )
    subprocess.check_call(["sudo", "umount", "-l", "/tmp/{}".format(_uuid)])
    _used = int(_res[2])

    return _used


def get_uuid(partition):
    logger.info("getting UUID for partition '{}'".format(partition))
    return (
        subprocess.check_output(["lsblk", "-no", "UUID", partition])
        .decode("utf-8")
        .strip()
    )


def label_partition(partition, label, fs=None):
    logger.info("labeling partition '{}' with label '{}'".format(partition, label))

    if fs is None:
        fs = (
            subprocess.check_output(["lsblk", "-no", "FSTYPE", partition])
            .decode("utf-8")
            .strip()
        )

    if fs == "btrfs":
        subprocess.check_call(
            ["sudo", "btrfs", "filesystem", "label", partition, label]
        )
    elif fs == "ext4":
        subprocess.check_call(["sudo", "e2label", partition, label])
    elif fs == "vfat":
        subprocess.check_call(["sudo", "fatlabel", partition, label])
    else:
        raise Exception("Unsupported filesystem '{}'".format(fs))

    return True


def umount_if(mountpoint):
    logger.info("unmounting '{}' if mounted".format(mountpoint))

    if os.path.ismount(mountpoint):
        subprocess.check_call(["sudo", "umount", "-l", mountpoint])


def remove_uuid_from_fstab(root, uuid):
    logger.info("removing UUID '{}' from fstab".format(uuid))
    subprocess.check_call(
        ["sudo", "sed", "-i", "/UUID={}/d".format(uuid), root + "/etc/fstab"]
    )

def add_bind_mount_to_fstab(src, dest, root):
    fstab_path = root + "/etc/fstab"
    bind_mount = "{} {} none bind 0 0\n".format(src, dest)
    with open(fstab_path, "a") as f:
        f.write(bind_mount)

def add_bind_mount_to_fstab(src, dest, root):
    fstab_path = root + "/etc/fstab"
    bind_mount = "{} {} none bind 0 0\n".format(src, dest)
    with open(fstab_path, "a") as f:
        f.write(bind_mount)


def update_grub(root, block_device, boot_partition, efi_partition=None):
    logger.info("updating GRUB in '{}'".format(root))

    umount_if(boot_partition)
    if efi_partition:
        umount_if(efi_partition)

    subprocess.check_call(["sudo", "mount", boot_partition, root + "/boot"])
    if efi_partition:
        subprocess.check_call(["sudo", "mount", efi_partition, root + "/boot/efi"])
    subprocess.check_call(["sudo", "mount", "--bind", "/dev", root + "/dev"])
    subprocess.check_call(["sudo", "mount", "--bind", "/dev/pts", root + "/dev/pts"])
    subprocess.check_call(["sudo", "mount", "--bind", "/proc", root + "/proc"])
    subprocess.check_call(["sudo", "mount", "--bind", "/sys", root + "/sys"])
    subprocess.check_call(["sudo", "mount", "--bind", "/run", root + "/run"])

    script = [
        "#!/bin/bash",
        "sudo chroot {} grub-mkconfig -o /boot/grub/grub.cfg".format(root),
    ]
    subprocess.check_call("\n".join(script), shell=True)

    subprocess.check_call(
        [
            "sudo",
            "grub-install",
            "--boot-directory",
            root + "/boot",
            "--target={}".format("x86_64-efi" if efi_partition else "i386-pc"),
            block_device,
        ]
    )
    script = [  # for some reason, grub-install doesn't work if we don't install it from the chroot too
        "#!/bin/bash",
        "sudo chroot {} grub-install --boot-directory /boot {} --target={}".format(
            root, block_device, "x86_64-efi" if efi_partition else "i386-pc"
        ),
    ]
    subprocess.check_call("\n".join(script), shell=True)

    umount_if(boot_partition)
    if efi_partition:
        umount_if(efi_partition)


def generate_grub_file(boot_uuid, root_a_uuid, root_b_uuid, kernel):
    boot_content = """#!/bin/sh
exec tail -n +3 $0

set menu_color_normal=white/black
set menu_color_highlight=black/light-gray

function gfxmode {
set gfxpayload="${1}"
if [ "${1}" = "keep" ]; then
        set vt_handoff=vt.handoff=7
else
        set vt_handoff=
fi
}
if [ "${recordfail}" != 1 ]; then
if [ -e ${prefix}/gfxblacklist.txt ]; then
    if [ ${grub_platform} != pc ]; then
    set linux_gfx_mode=keep
    elif hwmatch ${prefix}/gfxblacklist.txt 3; then
    if [ ${match} = 0 ]; then
        set linux_gfx_mode=keep
    else
        set linux_gfx_mode=text
    fi
    else
    set linux_gfx_mode=text
    fi
else
    set linux_gfx_mode=keep
fi
else
set linux_gfx_mode=text
fi
export linux_gfx_mode
"""
    boot_entry = """menuentry 'State %s' --class gnu-linux --class gnu --class os {
recordfail
load_video
gfxmode $linux_gfx_mode
insmod gzio
if [ x$grub_platform = xxen ]; then insmod xzio; insmod lzopio; fi
insmod part_gpt
insmod ext2
search --no-floppy --fs-uuid --set=root %s
linux	/vmlinuz-%s root=UUID=%s quiet splash bgrt_disable $vt_handoff
initrd  /initrd.img-%s
}
"""
    boot_content += boot_entry % ("A", boot_uuid, kernel, root_a_uuid, kernel)
    boot_content += boot_entry % ("B", boot_uuid, kernel, root_b_uuid, kernel)

    with open("/tmp/10_vanilla", "w") as f:
        f.write(boot_content)

    subprocess.check_call(["sudo", "chmod", "777", "/tmp/10_vanilla"])


def get_kernel_version(root):
    return sorted(os.listdir(root + "/usr/lib/modules"))[-1]


def entry_point(disk: dict, timezone: dict):
    logger.info("starting post_install script")

    block_device = disk["disk"]["disk"]
    tz_region = timezone["timezone"]["region"]
    tz_zone = timezone["timezone"]["zone"]

    root_partitions = []
    boot_partition = ""
    efi_partition = None
    bios_partition = None
    ignore_keys = ["disk", "auto"]
    for partition, info in disk["disk"].items():
        if partition not in ignore_keys:
            if info["mp"] == "/":
                root_partitions.append(partition)
            elif info["mp"] == "/boot":
                boot_partition = partition
            elif info["mp"] == "/boot/efi":
                efi_partition = partition
            elif info["mp"] == "":
                bios_partition = partition

    assert (
        len(root_partitions) == 2
    ), "Wrong number of root partitions provided. Expected 2"
    partition_sizes = [find_partition_used_space(part) for part in root_partitions]

    # Replace EFI grub with BIOS one because we need it for grub-install later
    if not efi_partition:
        subprocess.check_call(["sudo", "apt", "install", "-y", "grub-pc"])

    # Root A is the one filled by the installer, thus the root with more used disk space
    root_a = ""
    root_b = ""
    if partition_sizes[0] > partition_sizes[1]:
        root_a = root_partitions[0]
        root_b = root_partitions[1]
    else:
        root_a = root_partitions[1]
        root_b = root_partitions[0]

    # getting UUIDs
    logger.info("getting UUIDs for boot and root partitions")
    boot_uuid = get_uuid(boot_partition)
    root_a_uuid = get_uuid(root_a)
    root_b_uuid = get_uuid(root_b)

    # preparing mountpoints
    logger.info("preparing mountpoints")
    umount_if("/mnt")
    umount_if("/mnt/a")
    umount_if("/mnt/b")
    umount_if(root_a)
    umount_if(root_b)
    subprocess.check_call(["sudo", "mkdir", "-p", "/mnt/a", "/mnt/b"])

    # labeling root partitions
    logger.info("labeling root partitions")
    label_partition(root_a, "a")
    label_partition(root_b, "b")

    # mounting root partitions
    logger.info("mounting root partitions")
    subprocess.check_call(["sudo", "mount", root_a, "/mnt/a"])
    subprocess.check_call(["sudo", "mount", root_b, "/mnt/b"])

    # set timezone (workaround for distinst generating a broken /etc/timezone)
    # I am not going to debug this deeper, because it is a distinst bug and
    # we will switch to another backend soon or write our own
    subprocess.check_call("sudo rm -f /mnt/a/etc/timezone", shell=True)
    subprocess.check_call(
        "sudo bash -c 'echo \"%s/%s\" > /mnt/a/etc/timezone'" % (tz_region, tz_zone),
        shell=True,
    )
    subprocess.check_call("sudo rm -f /mnt/a/etc/localtime", shell=True)
    subprocess.check_call(
        "sudo ln -s /usr/share/zoneinfo/%s/%s /mnt/a/etc/localtime"
        % (tz_region, tz_zone),
        shell=True,
    )

    # adapting A strucutre
    logger.info("adapting A structure")
    subprocess.check_call(["sudo", "mkdir", "-p", "/mnt/a/.system"])
    subprocess.check_call("sudo mv /mnt/a/* /mnt/a/.system/", shell=True)

    # creating standard folders in A
    logger.info("creating standard folders in A")
    standard_folders = [
        "boot",
        "dev",
        "home",
        "media",
        "mnt",
        "partFuture",
        "proc",
        "root",
        "run",
        "srv",
        "sys",
        "tmp",
    ]
    for item in standard_folders:
        subprocess.check_call(["sudo", "mkdir", "-p", "/mnt/a/" + item])

    # creating relative links
    relative_links = [
        "usr",
        "etc",
        "root",
        "usr/bin",
        "usr/lib",
        "usr/lib32",
        "usr/lib64",
        "usr/libx32",
        "usr/sbin",
    ]
    relative_system_links = [
        "dev",
        "proc",
        "run",
        "srv",
        "sys",
        "tmp",
        "media",
        "boot",
    ]
    logger.info("creating relative links")
    script_a = ["#!/bin/bash", "cd /mnt/a/"]
    for link in relative_links:
        script_a.append("sudo ln -rs .system/{} .".format(link))
    subprocess.check_call(["sudo", "bash", "-c", "\n".join(script_a)])

    script_a = ["#!/bin/bash", "cd /mnt/a/"]
    for link in relative_system_links:
        script_a.append("sudo rm -rf .system/{}".format(link))
        script_a.append("sudo ln -rs {} .system/".format(link))
    subprocess.check_call(["sudo", "bash", "-c", "\n".join(script_a)])

    # removing unwanted UUIDs from A fstab
    logger.info("removing unwanted UUIDs from A fstab")
    remove_uuid_from_fstab("/mnt/a", root_b_uuid)

    # adding bind mounts from A fstab
    logger.info("adding bind mounts from A fstab")
    add_bind_mount_to_fstab("/.system/var", "/var", "/mnt/a")
    add_bind_mount_to_fstab("/.system/opt", "/opt", "/mnt/a")

    # getting kernel version
    logger.info("getting kernel version")
    kernel = get_kernel_version("/mnt/a")

    # generating 10_vanilla grub file
    logger.info("generating 10_vanilla grub file")
    generate_grub_file(boot_uuid, root_a_uuid, root_b_uuid, kernel)

    # adapting A grub
    logger.info("adapting A grub")
    grub_default = """GRUB_DEFAULT=1
GRUB_TIMEOUT=0
GRUB_HIDDEN_TIMEOUT=2
GRUB_TIMEOUT_STYLE=hidden"""

    with open("/mnt/a/.system/etc/default/grub", "w") as f:
        f.write(grub_default)

    subprocess.check_call(
        ["sudo", "cp", "/tmp/10_vanilla", "/mnt/a/.system/etc/grub.d/10_vanilla"]
    )
    subprocess.check_call(["sudo", "rm", "/mnt/a/.system/etc/grub.d/10_linux"])
    subprocess.check_call(["sudo", "rm", "/mnt/a/.system/etc/grub.d/20_memtest86+"])

    # rsyncing A to B
    logger.info("rsyncing A to B")
    subprocess.check_call(
        "sudo rsync -avxHAX --numeric-ids --exclude='/boot' --exclude='/dev' --exclude='/home' --exclude='/media' --exclude='/mnt' --exclude='/partFuture' --exclude='/proc' --exclude='/root' --exclude='/run' --exclude='/srv' --exclude='/sys' --exclude='/tmp' /mnt/a/ /mnt/b/",
        shell=True,
    )

    # creating standard folders in B
    logger.info("creating standard folders in B")
    standard_folders = [
        "boot",
        "dev",
        "home",
        "media",
        "mnt",
        "partFuture",
        "proc",
        "root",
        "run",
        "srv",
        "sys",
        "tmp",
    ]
    for item in standard_folders:
        subprocess.check_call(["sudo", "mkdir", "-p", "/mnt/b/" + item])

    # updating B fstab
    logger.info("updating B fstab")
    subprocess.check_call(
        [
            "sudo",
            "sed",
            "-i",
            "s/UUID={}/UUID={}/g".format(root_a_uuid, root_b_uuid),
            "/mnt/b/.system/etc/fstab",
        ]
    )

    # load efi modules
    if efi_partition:
        logger.info("load efi modules")
        subprocess.run(
            ["sudo", "modprobe", "efivars"], stderr=subprocess.DEVNULL
        )  # tested, safe to ignore starting from kernel 6.0

    # updating grub for both root partitions
    logger.info("updating grub for both root partitions")
    update_grub("/mnt/a", block_device, boot_partition, efi_partition)
    update_grub("/mnt/b", block_device, boot_partition, efi_partition)

    sys.exit(0)


if __name__ == "__main__":
    logger = logging.getLogger("Installer::PostInstall")
    logger.setLevel(logging.INFO)

    parser = argparse.ArgumentParser()
    parser.add_argument("disk", help="finals element containing disk information")
    parser.add_argument(
        "timezone", help="finals element containing timezone information"
    )
    args = parser.parse_args()

    disk_json = json.loads(args.disk)
    timezone_json = json.loads(args.timezone)

    try:
        entry_point(disk_json, timezone_json)
    except Exception as e:
        logger.error("Exception: {}".format(e))
        sys.exit(1)