mkosi.finalize

#!/usr/bin/env python3
# SPDX-License-Identifier: LGPL-2.1-or-later

import ast
import fnmatch
import functools
import os
import pathlib
import re
import shutil
import subprocess
import sys
import xattr

def dictify(f):
    def wrapper(*args, **kwargs):
        return dict(f(*args, **kwargs))
    return functools.update_wrapper(wrapper, f)

def buildroot():
    return pathlib.Path(os.getenv('BUILDROOT'))

def copy_path(oldpath, newpath, should_skip):
    newpath.mkdir(exist_ok=True)

    for entry in os.scandir(oldpath):
        if should_skip(entry.path):
            print(f'Not copying {entry.path}')

        newentry = newpath / entry.name
        if entry.is_dir(follow_symlinks=False):
            copy_path(entry.path, newentry, should_skip)
        elif entry.is_symlink():
            target = os.readlink(entry.path)
            newentry.symlink_to(target)
            shutil.copystat(entry.path, newentry, follow_symlinks=False)
        else:
            shutil.copy2(entry.path, newentry, follow_symlinks=False)

    shutil.copystat(oldpath, newpath, follow_symlinks=True)

def kernel_core_skip_file(path):
    patterns = ('*/bls.conf',
                '*/vmlinuz',
                '*/doc/*',
                '/boot/*')

    return any(fnmatch.fnmatch(path, pat)
               for pat in patterns)

def copy_in_modules_from_rpmls(root, kver, modules_rpm):
    rpm = f'{modules_rpm}-{kver}'
    try:
        out = subprocess.check_output(['rpm', '-ql', rpm], text=True)
    except (FileNotFoundError, subprocess.CalledProcessError) as e:
        print(f'Cannot query file list of {rpm}:', e)
        return False

    files = out.splitlines()
    installed, skipped = 0, 0
    for file in files:
        if kernel_core_skip_file(file):
            skipped += 1
            continue
        if pathlib.Path(file).is_dir():
            continue

        dest = root / file[1:]
        dest.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(file, dest, follow_symlinks=False)
        installed += 1

    print(f'Installed {installed}/{installed + skipped} files from {rpm}')
    return True

def copy_in_modules_from_fs(root, kver):
    source = f'/usr/lib/modules/{kver}'
    if not os.path.exists(source + '/kernel'):
        return False

    copy_path(pathlib.Path(source),
              root / source[1:],
              kernel_core_skip_file)
    print(f'Copied files from {source}')
    return True

def copy_in_modules_from_cpio(root, kver, modules_rpm):
    file_filter = ['--nonmatching',
                   './lib/modules/*/bls.conf',      # a grub abomination with $grub_variables
                   './lib/modules/*/vmlinuz']

    path = pathlib.Path(sys.argv[0]).parent
    file = path / f'{modules_rpm}-{kver}.rpm'

    if not file.exists():
        print(f'{file} not found, downloading.')

        subprocess.run(['dnf', 'download', f'--downloaddir={file.parent}',
                        file.with_suffix('').name],
                       check=True)

    with subprocess.Popen(['rpm2cpio', file], stdout=subprocess.PIPE) as archive:
        subprocess.run(['cpio', '-i', '--make-directories', '--quiet', '-D', root, *file_filter],
                       stdin=archive.stdout,
                       check=True)

    print(f'Unpacked files from {file}')
    return True

def copy_in_modules_kver(root, kver):
    print(f'Installing modules for kernel {kver}')

    res = subprocess.check_output(['rpm', '--eval', f'%[v"{kver}" >= v"6.2.0"]'],
                                  text=True)
    good = res.strip() == '1'
    modules_rpm = 'kernel-modules-core' if good else 'kernel-core'

    copy_in_modules_from_rpmls(root, kver, modules_rpm) or \
        copy_in_modules_from_fs(root, kver) or \
        copy_in_modules_from_cpio(root, kver, modules_rpm)

    subprocess.run(['depmod', '-a', '-w', '-b', root, kver], check=True)

def copy_in_modules(root):
    # This part is pretty ugly. The kernel package provides the following subpackages:
    # kernel — empty metapackage
    # kernel-core — /boot/vmlinuz-<kver> (%ghost?) and /lib/modules/<kver>/vmlinuz,
    # kernel-modules-core — a bunch of basic modules
    #        (k-m-c was split out in 6.2.0 out of kernel-core.)
    # kernel-modules — more modules (drivers)
    # kernel-modules-extra — more modules ("less commonly used drivers")
    # kernel-modules-internal — "kernel modules for the kernel package for Red Hat internal usage"
    #          (netdevsim, rcutorture…)
    #
    # This script requires kernel-[modules-]core)-<kver>.rpm to be present in $CWD.

    try:
        kver = os.environ['KERNEL_VERSION']
    except KeyError:
        kver = os.uname().release
        print(f'$KERNEL_VERSION not defined, using {kver}')

    copy_in_modules_kver(root, kver)

@dictify
def read_os_release(root):
    try:
        f = root.joinpath('etc/os-release').open()
    except FileNotFoundError:
        f = root.joinpath('usr/lib/os-release').open()

    for line_number, line in enumerate(f, start=1):
        if not line.strip() or line.startswith('#'):
            continue
        if m := re.match(r'([A-Z][A-Z_0-9]+)=(.*)', line):
            name, val = m.groups()
            if val and val[0] in '"\'':
                val = ast.literal_eval(val)
            yield name, val
        else:
            print(f'Bad line {line_number}: {line}', file=sys.stderr)

def update_suffixed(items, name, fallback):
    value = items.get(name, fallback)
    if 'mkosi-initrd' in value:
        return
    items[name] = value + ' (mkosi-initrd)'

def make_initrd_release(root, out):
    os_release = read_os_release(root)

    # Replacing fields in the original dictionary should maintain the order
    update_suffixed(os_release, 'NAME', 'Linux')
    update_suffixed(os_release, 'PRETTY_NAME', 'Linux')
    os_release['VARIANT'] = 'mkosi-initrd'
    os_release['VARIANT_ID'] = 'mkosi-initrd'

    for name, value in os_release.items():
        if value:
            print(f'{name}={value!r}', file=out)

    print(f'Writing {out.name} with PRETTY_NAME={os_release["PRETTY_NAME"]!r}')

def make_sysext_release(root, sysext_name, out):
    os_release = read_os_release(root)

    sysext = {
        'NAME': sysext_name,
        'SYSEXT_SCOPE': 'initrd',
    }
    for field in ('ID',
                  'VERSION_ID',
                  'VERSION_CODENAME',
                  'PLATFORM_ID',
                  'SYSEXT_LEVEL',

                  'HOME_URL',
                  'DOCUMENTATION_URL',
                  'SUPPORT_URL',
                  'BUG_REPORT_URL',
                  'PRIVACY_POLICY_URL'):
        if value := os_release.get(field):
            sysext[field] = value

    for name, value in sysext.items():
        print(f'{name}={value!r}', file=out)

    print(f'Writing {out.name} with NAME={sysext["NAME"]!r}')

def write_initrd_release(root):
    output = root / 'etc/initrd-release'
    output.unlink(missing_ok=True)

    with output.open('wt') as out:
        make_initrd_release(root, out)

    root.joinpath('etc/os-release').unlink(missing_ok=True)
    root.joinpath('usr/lib/os-release').unlink(missing_ok=True)

    root.joinpath('etc/os-release').symlink_to(output.name)

def write_sysext_release(root, sysext_name):
    output = root / f"usr/lib/extension-release.d/extension-release.{sysext_name}"
    output.parent.mkdir(exist_ok=True)

    with output.open('wt') as out:
        make_sysext_release(root, sysext_name, out)

    print('Setting user.extension-release.strict=0 on the extension-release file')
    xattr.set(output, 'user.extension-release.strict', '0')

def make_init_symlink(root):
    init = root / 'init'
    init.unlink(missing_ok=True)
    init.symlink_to('usr/lib/systemd/systemd')
    print(f'Symlinked {init} → usr/lib/systemd/systemd')

def make_sysroot_dir(root):
    sysroot = root / 'sysroot'
    sysroot.mkdir(mode=0o755, exist_ok=True)
    print(f'Created {sysroot}')

def make_debug_shell_emergency(root):
    unit = root / 'etc/systemd/system/emergency.service'
    unit.unlink(missing_ok=True)
    unit.symlink_to('debug-shell.service')
    print(f'Symlinked {unit} → debug-shell.service')

def mask_units(root):
    units = [
        # Dracut installs a rule which sets OPTIONS+="db_persist" on all dm devices [1]
        # (According to codesearch.debian.net, it is the only user of db_persist.)
        # Without this, all dm units end up with SYSTEMD_READY=0 and systemd thinks the
        # device units are missing. The system boots fine, but an attempt to call
        # daemon-reexec or daemon-reload ends with anything that can be stopped or unmounted
        # being purged.
        #
        # In systemd, we always have cleaned the database on switch-root, since the initial
        # addition of the initrd-* units in cf843477946451fabf9b5d17eec8ec81515057b6. But
        # that seems pointless, since initrds need to match the kernel version and are
        # generally used with the main system in the same version or slightly newer. And
        # for important devices, db_persist is set. So we end up destroying some of the
        # state, but not all. Let's just skip the cleanup altogether, and rely on the rules
        # being idempotent so that we end up in the correct state.
        #
        # [1] https://raw.githubusercontent.com/dracutdevs/dracut/2d83bce21bfc874b29c1fb99e8fabb843f038725/modules.d/90dm/11-dm.rules
        'initrd-udevadm-cleanup-db.service',
    ]

    subprocess.run(['systemctl', f'--root={root}', 'mask', *units], check=True)

def remove_non_sysext_files(root):
    var = root / "var"
    if var.exists():
        shutil.rmtree(var)

    etc = root / "etc"
    if etc.exists():
        # from mkosi import find_files
        # for path in find_files(etc):
        #     print(f'WARNING: configuration file: {path}')
        # … This doesn't work because we also see files from the lower layer

        shutil.rmtree(etc)


def do_initrd(root):
    copy_in_modules(root)

    write_initrd_release(root)

    make_init_symlink(root)
    make_sysroot_dir(root)

    make_debug_shell_emergency(root)

    mask_units(root)

def do_sysext(root, sysext_name):
    write_sysext_release(root, sysext_name)
    remove_non_sysext_files(root)

def main(argv):
    if argv[1] != 'final':
        sys.exit(0)

    root = buildroot()

    sysext = os.getenv('SYSEXT')
    if sysext:
        do_sysext(root, sysext)
    else:
        do_initrd(root)

if __name__ == '__main__':
    main(sys.argv)