create-initrd

#!/usr/bin/env python3
import argparse
import os

from config import CONFIG as cfg
from util import (
    command,
    download_tar,
    ensure_root,
    get_consent,
    initrd_write,
    list_files_in_packages,
    mount_tmpfs,
    umount,
    warn,
)

ensure_root()

cli = argparse.ArgumentParser()
cli.add_argument('--debian-mirror', default='http://mirror.stusta.de/debian/')
cli.add_argument('--out', default=cfg.path.cpio, help="output cpio file (%(default)s)")
cli.add_argument('--compressor', default=cfg.packing.compressor)
cli.add_argument('--skip-setup', action='store_true')
cli.add_argument('--update', action='store_true')
cli.add_argument('--tmp-work', action='store_true',
                 help='use a tmpfs for work and cachedir')
args = cli.parse_args()

# Some random sanity checks
selected_system_modules = sum(
    [mod in ('system-debian', 'system-arch', 'system-arch-dracut', 'system-gentoo') for mod in cfg.modules])

if selected_system_modules != 1:
    warn("Please select exactly one system-* module")
    raise SystemExit(1)

if command('uname', '-r', capture_stdout=True).decode()[:-1] not in os.listdir('/lib/modules'):
    warn('Modules for running kernel not found in /lib/modules. If you just updated your arch linux, you probably '
         'should reboot before continuing.\nIf you are not running arch, I don\'t know how you triggered this '
         'message.\nDo you want to continue?')
    if not get_consent():
        raise SystemExit(1)

if not args.skip_setup:
    # discard existing ramdisk content
    umount(cfg.path.initrd_devel)
    umount(cfg.path.work)

    os.makedirs(cfg.path.work, exist_ok=True)
    os.makedirs(cfg.path.cache, exist_ok=True)

    # provide the various tmpfses
    if args.tmp_work:
        mount_tmpfs(cfg.path.work)
        mount_tmpfs(cfg.path.cache)

    deb_packages = [
        'ifrename',  # needed by the payload scripts
        'iproute2',  # needed by the payload scripts
        'kexec-tools',  # needed for booting the payload system
        'linux-image-amd64',  # needed for booting the stiefel system
        'python3',  # all payload scripts are written in Python3
        'python3-pycryptodome',
        'python3-aiohttp',
        'systemd-container',  # to allow launching inside nspawn during this script
        'dbus-broker',  # apparently required for systemd-container
        'systemd-sysv',  # to provide symlinks in /sbin: init, poweroff, ...
    ]

    if 'nbd' in cfg.modules:
        deb_packages.extend(['nbd-server'])

    if 'lvm' in cfg.modules:
        deb_packages.append('lvm2')

    if 'i915' in cfg.modules:
        deb_packages.append('firmware-misc-nonfree')

    if cfg.boot.luks_block:
        deb_packages.append('cryptsetup')

    deb_packages.extend(cfg.initrd.include_packages)

    # perform initial setup
    # this logs to $workdir/$workdir_subpath_initrd/debootstrap/debootstrap.log
    # -> usually "workdir/initrd.nspawn/debootstrap/debootstrap.log"
    print(f"running debootstrap with log {cfg.path.initrd}/debootstrap/debootstrap.log")
    command(
        'debootstrap',
        '--include=' + ','.join(deb_packages),
        '--cache-dir=' + os.path.abspath(cfg.path.cache),
        '--variant=minbase',
        '--components=main,contrib,non-free',
        '--merged-usr',
        '--verbose',
        '--merged-usr',
        'stable',
        cfg.path.initrd,
        args.debian_mirror
    )

# install our overlay
command('cp', '-RT', 'overlays/initrd', cfg.path.initrd)

# install the AES key
if not os.path.exists('aes-key'):
    print('generating new AES key')
    with open('aes-key', 'wb') as fileobj:
        fileobj.write(os.urandom(16))
command('cp', 'aes-key', cfg.path.initrd)

if not args.skip_setup:
    # set root password
    command(
        'chpasswd',
        nspawn=cfg.path.initrd,
        stdin=f'root:{cfg.initrd.password}\n'
    )

    # set root shell
    command(
        'chsh', '-s',
        cfg.initrd.shell,
        nspawn=cfg.path.initrd
    )

    # enable login from systemd-nspawn
    initrd_write('/etc/securetty', 'pts/0', append=True)

    # set the hostname
    initrd_write('/etc/hostname', 'stiefelsystem')

    # disable lidswitch handling...
    initrd_write('/etc/systemd/logind.conf',
        'HandleSuspendKey=ignore',
        'HandleHibernateKey=ignore',
        'HandleLidSwitch=ignore',
        'HandleLidSwitchExternalPower=ignore',
        'HandleLidSwitchDocked=ignore',
        append=True
    )

    # systemd configuration
    command('systemctl', 'set-default',
        'multi-user.target',
        nspawn=cfg.path.initrd
    )
    command('systemctl', 'enable',
        'fake-entropy.service',
        nspawn=cfg.path.initrd
    )
    command('systemctl', 'disable',
        'kexec.service',
        'kexec-load.service',
        'machines.target',
        'remote-fs.target',
        nspawn=cfg.path.initrd
    )
    command('systemctl', 'mask',
        'serial-getty@.service',
        'apt-daily.timer',
        'apt-daily-upgrade.timer',
        'serial-getty@.service',
        'systemd-journal-flush.service',
        'systemd-timedated.service',
        'systemd-timesyncd.service',
        'systemd-tmpfiles-clean.timer',
        'systemd-update-utmp.service',
        'systemd-update-utmp-runlevel.service',
        'time-sync.target',
        nspawn=cfg.path.initrd
    )

    # stiefel configuration
    if 'nbd' in cfg.modules:
        command('systemctl', 'disable', 'nbd-server.service', nspawn=cfg.path.initrd)
        # the nbd server configuration will be generated on-the-fly

    # kernel name
    kernel_name = os.readlink(cfg.path.initrd + '/vmlinuz')[13:]

    # create devel overlay system which we'll use to compile a few things
    os.makedirs(cfg.path.initrd_devel, exist_ok=True)
    os.makedirs(cfg.path.initrd_devel + "-overlayfs-upperdir", exist_ok=True)
    os.makedirs(cfg.path.initrd_devel + "-overlayfs-workdir", exist_ok=True)

    command(
        "mount",
        "-t", "overlay",
        "overlay",
        "-o", ",".join([
            f"lowerdir={os.path.abspath(cfg.path.initrd)}",
            f"upperdir={os.path.abspath(cfg.path.initrd_devel + '-overlayfs-upperdir')}",
            f"workdir={os.path.abspath(cfg.path.initrd_devel + '-overlayfs-workdir')}",
        ]),
        os.path.abspath(cfg.path.initrd_devel),
    )

    command('apt', 'update', nspawn=cfg.path.initrd_devel)
    command('apt', 'upgrade', nspawn=cfg.path.initrd_devel)
    command('apt', 'install', '-y',
        'build-essential',
        'git',
        'linux-headers-amd64',
        'kmod',
        nspawn=cfg.path.initrd_devel
    )

    # compile and install  the userland driver for the clevo fan controller
    if 'clevo-fancontrol' in cfg.modules:
        download_tar(
            cfg.mod_config["clevo-fancontrol"].url,
            os.path.join(cfg.path.initrd_devel, "root/fancontrol")
        )
        command('make', '-C', '/root/fancontrol', nspawn=cfg.path.initrd_devel)
        command(
            "cp",
            cfg.path.initrd_devel + "/root/fancontrol/clevo-fancontrol",
            cfg.path.initrd + "/usr/local/bin",
        )
        command(
            "cp",
            cfg.path.initrd_devel + "/root/fancontrol/clevo-fancontrol.service",
            cfg.path.initrd + "/etc/systemd/system",
        )
        command('systemctl', 'enable',
            'clevo-fancontrol.service',
            nspawn=cfg.path.initrd
        )

    if 'r8152' in cfg.modules:
        download_tar(
            cfg.mod_config["r8152"].url,
            os.path.join(cfg.path.initrd_devel, "root/r8152")
        )

        command(
            'make',
            '-C', '/root/r8152',
            f'KERNELDIR=/lib/modules/{kernel_name}/build',
            'modules',
            nspawn=cfg.path.initrd_devel
        )
        command(
            "cp",
            cfg.path.initrd_devel + "/root/r8152/r8152.ko",
            cfg.path.initrd + f"/lib/modules/{kernel_name}/kernel/drivers/net/usb/",
        )
        os.makedirs(cfg.path.initrd + "/etc/udev/rules.d", exist_ok=True)
        command(
            "cp",
            cfg.path.initrd_devel + "/root/r8152/50-usb-realtek-net.rules",
            cfg.path.initrd + "/etc/udev/rules.d/"
        )
        command(
            "depmod",
            "-a",
            kernel_name,
            nspawn=cfg.path.initrd,
        )

if args.update:
    command('apt', 'update', nspawn=cfg.path.initrd_devel)
    command('apt', 'upgrade', nspawn=cfg.path.initrd_devel)

if 'debug' in cfg.modules:
    command('updatedb', nspawn=cfg.path.initrd)

if not args.skip_setup:
    if 'debug' not in cfg.modules:
        # strip kernel modules
        count = 0
        for path, _, files in os.walk(cfg.path.initrd + '/lib/modules'):
            for filename in files:
                if filename.endswith('.ko'):
                    full_path = os.path.join(path, filename)
                    command('strip', '--strip-debug', full_path, silent=count)
                    count += 1
        if count > 1:
            print(f'$ ... (stripped {count - 1} more modules)')

# the folder is ready, now we can pack and compress the initrd

paths_to_exclude = set(
    list_files_in_packages(cfg.packing.exclude_packages, cfg.path.initrd)
)
paths_to_exclude.update(
    path.encode() for path in cfg.packing.exclude_paths
)

old_cwd = os.getcwd()
os.chdir(cfg.path.initrd)

target_files = []

def scan_path(path):
    for name in os.listdir(path):
        full = os.path.normpath(os.path.join(path, name))

        if full in paths_to_exclude:
            continue
        if full.endswith(b'__pycache__'):
            continue

        yield full
        # recurse into directories (don't follow links)
        if os.path.isdir(full) and not os.path.islink(full):
            yield from scan_path(full)

print(f'packing {cfg.path.initrd}')

archive = command(
    "bsdcpio", "-0", "-o", "-H", "newc",
    stdin=b'\0'.join(scan_path(b'.')) + b'\0',
    capture_stdout=True,
    env={"LANG": "C"},
)

os.chdir(old_cwd)
del old_cwd

print(f"uncompressed CPIO: {len(archive)} bytes")

compressed = command(
    f"pv -s {len(archive)} | {args.compressor}",
    shell=True,
    stdin=archive,
    capture_stdout=True
)
print(f"compressed CPIO: {len(compressed)} bytes")

with open(args.out, "wb") as cpiofile:
    cpiofile.write(compressed)