#!python
#
# Script building an initramfs
# Reads configuration file from :
#   - CMKINITCFG file if environment variable set
#   - /etc/cmkinitramfs.ini by default
# The CMKINIT_SCRIPT environment variable override the path of the
# cmkinit script, which is by default searched in PATH
#

import cmkinitramfs
import configparser
import os
import argparse
import sys
import subprocess
import shutil
import glob

def info(text):
    """Print information string for the user"""
    if not args.quiet:
        print(text, file=sys.stderr)

# Parse command line
parser = argparse.ArgumentParser(description="Build an initramfs.")
parser.add_argument(
    "--debug", "-d", dest="debug", action="store_true",
    help="Enable debugging mode: non root and does not create final archive"
)
parser.add_argument(
    "--dry-run", "-D", dest="dryrun", action="store_true",
    help="Enable dry-run: does not creates final initramfs"
)
parser.add_argument(
    "--output", "-o", dest="output", type=str,
    help="Set output cpio file and disable writing to /boot"
)
parser.add_argument(
    "--clean", "-C", dest="clean", action="store_true",
    help="Overwrite temporary directory if it exists"
)
parser.add_argument(
    "--quiet", "-q", dest="quiet", action="store_true",
    help="Don't output status informations"
)
parser.add_argument(
    "kernel", type=str, nargs='?',
    help="Select kernel version to cleanup rather than /usr/src/linux"
)
args = parser.parse_args()

# Load configuration
config = configparser.ConfigParser()
if os.environ.get("CMKINITCFG"):
    config.read(os.environ["CMKINITCFG"])
else:
    config.read("/etc/cmkinitramfs.ini")

# Configure initramfs module
if config["DEFAULT"].get("build-dir"):
    initramfs.DESTDIR = config["DEFAULT"]["build-dir"].strip()
cmkinitramfs.QUIET = args.quiet

# types: Set of data types used to boot
types = set()
for data in [k for k in config if k != "DEFAULT"]:
    types.add(config[data]["type"])

# filesystems: Set of filesystems used to boot
filesystems = set()
for mount in [k for k in config if k != "DEFAULT" \
              and config[k]["type"] == "mount"]:
    filesystems.add(config[mount]["filesystem"])

# Cleanup

if args.clean:
    info(f"Warning: Overwriting temporary directory {cmkinitramfs.DESTDIR}")
    cmkinitramfs.cleanup()

# Build initramfs

info("Building initramfs")
cmkinitramfs.mklayout(debug=args.debug)

# Copy user files
if config["DEFAULT"].get("files"):
    for filepath in config["DEFAULT"]["files"].strip().split(':'):
        info(f"Copying {filepath} to /root")
        cmkinitramfs.copyfile(filepath)

# Busybox
info("Installing busybox")
cmkinitramfs.copyexec("busybox")
cmkinitramfs.install_busybox()

# Cryptsetup
if "luks" in types:
    info("Installing LUKS utils")
    cmkinitramfs.copyexec("cryptsetup")
    cmkinitramfs.copylib("libgcc_s.so.1")

# LVM
if "lvm" in types:
    info("Installing LVM utils")
    cmkinitramfs.copyexec("lvm")

# md
if "md" in types:
    info("Installing MD utils")
    cmkinitramfs.copyexec("mdadm")

# BTRFS
if "btrfs" in filesystems:
    info("Installing BTRFS utils")
    cmkinitramfs.copyexec("btrfs")
    cmkinitramfs.copyexec("fsck.btrfs")

# EXT4
if "ext4" in filesystems:
    info("Installing EXT4 utils")
    cmkinitramfs.copyexec("fsck.ext4")
    cmkinitramfs.copyexec("e2fsck")

# XFS
if "xfs" in filesystems:
    info("Installing XFS utils")
    cmkinitramfs.copyexec("fsck.xfs")
    cmkinitramfs.copyexec("xfs_repair")

# FAT
if "fat" in filesystems or "vfat" in filesystems:
    info("Installing FAT utils")
    cmkinitramfs.copyexec("fsck.fat")
    cmkinitramfs.copyexec("fsck.vfat")

# F2FS
if "f2fs" in filesystems:
    info("Installing F2FS utils")
    cmkinitramfs.copyexec("fsck.f2fs")

# Keyboard
if config["DEFAULT"].get("keymap"):
    info("Copying keymap to /root")
    gzip = subprocess.run(
        ["gzip", "-cd", config["DEFAULT"]["keymap"].strip()],
        stdout=subprocess.PIPE, check=True
    )
    loadkeys = subprocess.run(
        ["loadkeys", "--bkeymap"],
        input=gzip.stdout, stdout=subprocess.PIPE, check=True
    )
    cmkinitramfs.writefile(
        loadkeys.stdout,
        config["DEFAULT"].get("keymap-file", "/root/keymap.bmap")
    )

# Create /init
info("Generating /init")
init = subprocess.run([os.environ.get("CMKINIT_SCRIPT", "cmkinit")],
    stdout=subprocess.PIPE, check=True)
cmkinitramfs.writefile(init.stdout, "/init", 0o755)

# Find and hardlink duplicated files
info("Hardlinking duplicated files")
cmkinitramfs.hardlink_duplicates()

# Create initramfs
if not(args.debug) and not(args.dryrun):
    # Build cpio
    info("Building CPIO archive")
    cpio = cmkinitramfs.mkcpio()
    if args.output == "-":
        sys.stdout.buffer.write(cpio)
    else:
        output = args.output if args.output else "/usr/src/initramfs.cpio"
        with open(output, "wb") as filedest:
            filedest.write(cpio)

    # Cleanup temporary directory
    info("Cleaning up temporary files")
    cmkinitramfs.cleanup()

    # Cleanup kernel's initramfs
    if args.kernel != "none":
        kver = args.kernel if args.kernel else "linux"
        info(f"Cleaning up kernel {kver}")
        rm = glob.glob(f"/usr/src/{kver}/usr/initramfs_data.cpio*")
        for fname in rm:
            os.remove(fname)

    if not args.output:
        # Copy cpio to /boot
        info("Installing initramfs to /boot")
        if os.path.isfile("/boot/initramfs.cpio.xz"):
            shutil.copyfile("/boot/initramfs.cpio.xz",
                            "/boot/initramfs.cpio.xz.old")
        xz = subprocess.run(
            ["xz", "-zkc", "--check=crc32", "/usr/src/initramfs.cpio"],
            stdout=subprocess.PIPE, check=True
        )
        with open("/boot/initramfs.cpio.xz", "wb") as filedest:
            filedest.write(xz.stdout)

info("Success")

