initramfs-compression-test.py

#! /usr/bin/env python3

import os
from pathlib import PosixPath
import re
import subprocess


def main():
    if os.getuid() != 0:
        raise SystemExit(f"\n Root access is required to test initramfs compression.\n")

    determine_initramfs_compression()


def get_initramfs_paths():
    initramfs_paths = []
    ostree_booted = PosixPath("/run/ostree-booted")
    if ostree_booted.exists():
        print(f"\n Operating in image mode\n")
        entries_dir = PosixPath("/boot/loader/entries")
        for entry in entries_dir.iterdir():
            ostree_conf = entry.read_text()
            initrd_re_compile = re.compile(r"^initrd\b.+\.img.*", re.MULTILINE)
            initrd_line = re.search(initrd_re_compile, ostree_conf)
            initramfs_path_image_mode = ostree_conf[initrd_line.start():initrd_line.end():].replace("initrd ", "/boot")
            initramfs_paths.append(initramfs_path_image_mode)
    else:
        print(f"\n Operating in package mode\n")
        initramfs_paths_package_mode = list(PosixPath("/boot").glob("initramfs*"))
        for initramfs_path in initramfs_paths_package_mode:
            initramfs_paths.append(str(initramfs_path))

    return initramfs_paths


def determine_initramfs_compression():
    FILE_SIGNATURES = {
        "Gzip": r"\x1f\x8b\x08",
        "XZ": r"\xfd\x37\x7a\x58\x5a\x00",
        "Zstandard": r"\x28\xb5\x2f\xfd",
    }
    for initramfs_path in get_initramfs_paths():
        extract_initramfs = subprocess.run(
            ["/usr/lib/dracut/skipcpio", 
            initramfs_path], 
            capture_output=True, check=True)

        for compression, file_signature in FILE_SIGNATURES.items():
            if re.search(bytes(file_signature, encoding="UTF-8"), extract_initramfs.stdout[:8]):
                print(f" Initramfs: {PosixPath(initramfs_path)}")
                print(f" Compression: {compression}\n")


if __name__ == "__main__":
    raise SystemExit(main())