setup-server-os

#!/usr/bin/env python3
"""
script to prepare your OS for being stiefeled.
"""
import argparse
import hashlib
import json
import os
import shutil

from config import CONFIG as cfg
from util import (
    command,
    ensure_root,
    ensure_unit_enabled,
    restart_unit,
    FileEditor,
    install_folder,
)

ensure_root()
cli = argparse.ArgumentParser()
args = cli.parse_args()

# ensure that all required tools are installed
if shutil.which('ifrename') is None:
    raise RuntimeError("could not find ifrename")
try:
    import pyudev
    del pyudev
except ImportError:
    raise RuntimeError("could not find pyudev") from None


if 'base' in cfg.modules:
    # base configuration (for all distros)
    install_folder('overlays/server-os-generic')
    ensure_unit_enabled('stiefel-autokexec.service')

    with open('aes-key', 'rb') as keyfileobj:
        KEY = keyfileobj.read()

    stiefel_config = {
        'aes-key-hash': hashlib.sha256(KEY).hexdigest(),
        'hmac-key': hashlib.sha256(b'autokexec-reboot/' + KEY).hexdigest(),
        "autokexec-triggers": {
            "mac_detection": cfg.autokexec.macs,
            "broadcast": cfg.autokexec.broadcast,
            "adapters": cfg.autokexec.macs,
        },
        "bootdisk": cfg.boot.disk,
        "bootpart-luks": cfg.boot.luks_block,
        "bootpart": cfg.boot.part,
        "stiefelsystem-kernel": cfg.server_setup.stiefelsystem_kernel,
        "stiefelsystem-initrd": cfg.server_setup.stiefelsystem_initrd,
        "cmdline": cfg.server_setup.cmdline,
    }
    # TODO: rename to autokexec config
    edit = FileEditor('/etc/stiefelsystem/config.json')
    edit.set_data(json.dumps(stiefel_config, indent=4).encode() + b'\n')
    kexecd_config_changed = edit.write()

    if kexecd_config_changed:
        restart_unit('stiefel-autokexec.service')

    # copy the special stiefelsystem initrd to its final location
    # it will be used for autokexec on the server
    edit = FileEditor(cfg.server_setup.stiefelsystem_initrd)
    edit.load_from(os.path.join(cfg.path.work, 'initrd.cpio'))
    edit.write()

    # same for the debian-generic-kernel
    edit = FileEditor(cfg.server_setup.stiefelsystem_kernel)
    edit.load_from(os.path.join(cfg.path.initrd, 'vmlinuz'))
    edit.write()

if 'system-arch' in cfg.modules:
    # create a new 'stiefel' mkinitcpio-preset
    # which will generate a regular arch-initrd
    # with stiefel-client mounting support
    edit = FileEditor('/etc/mkinitcpio-stiefel.conf')
    edit.load_from('/etc/mkinitcpio.conf')
    changes = {'`which ifrename`': 'at-end'}
    edit.edit_bash_list('BINARIES', changes)
    changes = {'amdgpu': 'at-end', 'i915': 'at-end'}
    if 'r8152' in cfg.modules:
        changes['r8152'] = 'at-end'
    edit.edit_bash_list('MODULES', changes)
    changes = {'autodetect': 'remove'}
    if 'nbd' in cfg.modules:
        changes['nbd'] = 'before-fsck'
    edit.edit_bash_list('HOOKS', changes)
    edit.write()

    edit = FileEditor('/etc/mkinitcpio.d/linux.preset')
    edit.load()
    edit.edit_bash_list('PRESETS', {'stiefel': 'at-end'})
    edit.add_or_edit_var('stiefel_image', '/boot/initramfs-stiefel.img', add_prefix='\n')
    edit.add_or_edit_var('stiefel_options', '-c /etc/mkinitcpio-stiefel.conf -S autodetect')
    edit.write()

    # what kernel and initrd is then transferred and kexecd on the stiefel-client
    boot_config = {
        "kernel": "vmlinuz-linux",
        "initrd": "initramfs-stiefel.img",
        "cmdline": cfg.boot.cmdline,
        "stiefelmodules": list(cfg.modules.keys()),
    }
    edit = FileEditor('/boot/stiefelsystem.json')
    edit.set_data(json.dumps(boot_config, indent=4).encode() + b'\n')
    edit.write()

    if 'nbd' in cfg.modules:
        install_folder('overlays/server-os-arch-nbd')

    command('mkinitcpio', '-p', 'linux')

elif 'system-debian' in cfg.modules:
    if 'nbd' in cfg.modules:
        install_folder('overlays/server-os-debian')
        command('update-initramfs', '-u', '-k', 'all')

elif any(mod in ('system-gentoo', 'system-arch-dracut') for mod in cfg.modules):
    # tell stiefel-server which kernel to serve to a stiefel-client,
    # relative to /boot
    boot_config = {
        "kernel": cfg.boot.kernel,
        "initrd": cfg.boot.initrd,
        "cmdline": cfg.boot.cmdline,
        "stiefelmodules": list(cfg.modules.keys()),
    }
    edit = FileEditor('/boot/stiefelsystem.json')
    edit.set_data(json.dumps(boot_config, indent=4).encode() + b'\n')
    edit.write()

    if 'nbd' in cfg.modules:
        install_folder('overlays/server-os-dracut-nbd')
        if 'system-arch-dracut' in cfg.modules:
            install_folder('overlays/server-os-dracut-arch')
        command('dracut', cfg.boot.initrd, '--add', '" nbd "', '--no-hostonly', '--force')

else:
    raise Exception("no system-specific distro config module is enabled")