setup-client-usbdrive

#!/usr/bin/env python3
"""
script to create the usb flash drive for booting the client
"""
import argparse
import base64
import os
import tempfile

from config import CONFIG as cfg
from util import (
    command,
    ensure_root,
    get_consent,
    warn,
)

ensure_root()

command('lsblk')

cli = argparse.ArgumentParser()
cli.add_argument('blockdev')
# some older BIOSes only boot if the boot partition starts at sector 32...
cli.add_argument('--first-sector', type=int, default=32)
args = cli.parse_args()

if not os.path.exists(args.blockdev):
    cli.error(f'block device does not exist: {args.blockdev!r}')

partition = f'{args.blockdev}1'

warn(f'wiping entire drive at {args.blockdev!r} to create {partition!r}')

if not get_consent():
    raise SystemExit(1)

# TODO create gpt table

# create partition table and write MBR
command('sfdisk', args.blockdev, stdin=f"label: dos\n{args.first_sector},1638400,c,*\n")
if any(mod in ('system-arch', 'system-arch-dracut') for mod in cfg.modules):
    command('dd', 'if=/usr/lib/syslinux/bios/mbr.bin', 'of=' + args.blockdev)
elif 'system-debian' in cfg.modules:
    command('dd', 'if=/usr/lib/syslinux/mbr/mbr.bin', 'of=' + args.blockdev) #debian
elif 'system-gentoo' in cfg.modules:
    command('dd', 'if=/usr/share/syslinux/mbr.bin', 'of=' + args.blockdev)
else:
    print("no system specified in config.yaml modules")
    exit(1)

# create filesystem
command('mkfs.vfat', '-F', '16', partition)
# install bootloader
command('syslinux', partition)

# mount the filesystem and create the files on it
with tempfile.TemporaryDirectory() as tmpdir:
    command('mount', partition, tmpdir)
    try:
        command(
            'dd',
            'bs=1M',
            'if=' + os.path.join(cfg.path.work, 'initrd.cpio'),
            'of=' + os.path.join(tmpdir, 'initrd'),
            'oflag=direct',
            'status=progress',
        )
        command(
            'dd',
            'bs=1M',
            'if=' + os.path.join(cfg.path.initrd, 'vmlinuz'),
            'of=' + os.path.join(tmpdir, 'kernel'),
            'oflag=direct',
            'status=progress',
        )

        cmdline = " ".join([
            # the client system shouldn't modeset.
            # if it modesets, then the early boot steps of the actual target
            # initrd won't have working video output, making debugging
            # them harder.
            "nomodeset",
            "systemd.unit=stiefel-client.service",
        ])

        with open(os.path.join(tmpdir, 'syslinux.cfg'), 'w') as syslinuxcfg:
            syslinuxcfg.write(f"default kernel initrd=initrd {cmdline}\n")
    finally:
        command('umount', tmpdir)

print("synching io buffers...")
os.sync()
print("done.")