#!/bin/bash
# SPDX-License-Identifier: GPL-2.0-only
#
# vm-initrd: Create initramfs that just modprobe and switch to rootfs
#
# Copyright (C) 2023 Vitaly Chikunov <vt@altlinux.org>
#
# shellcheck disable=SC2155,SC2207,SC2128

set -eo pipefail

show_usage()
{
	echo "Usage: $0 [OPTIONS...] IMAGE [KERNEL_VERSION]"
	echo "Create minimalistic initramfs image"
	echo
	echo "  -b, --basedir=     Modules basedir (where /lib/modules located)."
	echo "  --modules=         Add modules into IMAGE."
	echo "  --add=X[:Y]        Add file/dir X into IMAGE at path Y."
	echo "  --list             List initramfs content."
	echo "  -v, --verbose      Increase verbosity level."
	echo "  -h, --help         This help."
	exit "$@"
}

warning() {
	local w prefix="$(basename "$0"): "
	for w do
		echo >&2 "$prefix$w"
		prefix=
	done
}

fatal() {
	warning "$@"
	exit 1
}

ROOT=
MODULES=
VERBOSE=
FILES=
LIST=

args=()
while [ $# -gt 0 ]; do # opt
	opt=$1
	arg=${opt##*=}
	case $opt in
		--basedir=*)	ROOT=$arg ;;
		-b)		shift; ROOT="${1?"$opt requires argument"}" ;;
		--modules=*)	MODULES+=" $arg" ;;
		--add=*)	FILES+=" $arg" ;;
		-l | --list)	LIST=y ;;
		-v | --verbose)	VERBOSE=1 ;;
		-h | --help)	show_usage ;;
		--)		shift; break ;;
		-*)		fatal "Unknown option $opt" ;;
		*)		args+=( "$opt" ) ;;
	esac
	shift
done
set -- "${args[@]}" "$@"

[ -n "${1-}" ] || { warning "Specifying IMAGE is required." ""; show_usage 1; }
INITRD="$1"; shift

if [ -n "$LIST" ]; then
	# Need to disable pipefail for dd(1).
	set +o pipefail
	count=0
	mime=$(file -b -L --mime-type "$INITRD")
	# List microcode.
	case "$mime" in
		application/x-cpio)
			{ count=$(cpio -tv < "$INITRD" 2>&1 >&3); } 3>&1
			count=${count% blocks}
			;;
	esac
	mime=$(dd if="$INITRD" skip="$count" 2>/dev/null | file -b --mime-type -)
	case "$mime" in
		application/x-cpio)			uncompress="cat" ;;
		application/x-gzip|application/gzip)	uncompress="gzip -dc" ;;
		application/x-xz)			uncompress="xz -dc" ;;
		application/x-zstd|application/zstd)	uncompress="zstd -dc" ;;
		*) fatal "Unknown file type $mime for $INITRD" >&2; ;;
	esac
	dd if="$INITRD" skip="$count" 2>/dev/null | $uncompress | cpio -tv --quiet
	exit
fi

if [ -n "$ROOT" ] &&
   [ ! -d "$ROOT/lib/modules" ] &&
   [ -d "$ROOT/usr/lib/modules" ]; then
	ROOT="$ROOT/usr"
fi

if [ -n "${1-}" ]; then
	KVER="$1"; shift
else
	KVER=( $(cd  "$ROOT/lib/modules" && ls) )
	[ ${#KVER[@]} -gt 1 ] && fatal "Too many kernels found (${#KVER[@]}), specify one."
	[ ${#KVER[@]} -eq 0 ] && fatal "No kernels found, install one or adjust --basedir="
fi
[ -d "$ROOT/lib/modules/$KVER" ] || fatal "Requested kernel missing $ROOT/lib/modules dir."

[ $# -eq 0 ] || fatal "Exceeding options: $*"

# Will continue generating initrd on file copying failures, but report failure (2) at exit.
RET=0

CPIO=$(mktemp -d)

# Pre-parse to catch dependencies of non-installed modules.
MODS=
for m in $MODULES; do
	if [ -n "${m#*.ko*}" ]; then
		# Installed module doesn't have '.ko' in the name.
		MODS+=" $m"
	else
		[ "${m#*/}" != "$m" ] || [ ! -e "$m" ] || m="./$m"
		# Non-installed module.
		depends=$(/sbin/modinfo -b "$ROOT" -k "$KVER" "$m" -F depends)
		MODS+=" $depends"
		softdep=$(/sbin/modinfo -b "$ROOT" -k "$KVER" "$m" -F softdep)
		pre=${softdep#*pre:}
		pre=${pre%post:*}
		post=${softdep#*post:}
		post=${post%pre:*}
		MODS+=" $pre $m $post"
	fi
done
MODULES=${MODS# }
unset MODS
true > "$CPIO/modprobe.txt"
# shellcheck disable=SC2086
if [ -n "$MODULES" ] &&
   ! /sbin/modprobe -d "$ROOT" -S "$KVER" -D -a $MODULES > "$CPIO/modprobe.txt"; then
	echo "initrd: Error: modprobe failure." >&2
	if [ ! -s "$ROOT/lib/modules/$KVER/modules.dep.bin" ]; then
		[ -w "$ROOT/lib/modules/$KVER/" ] && HELP= \
			|| HELP=" under rooter"
		echo "initrd:   Run 'depmod${ROOT:+ -b $ROOT} $KVER'$HELP." >&2
	fi
	exit 1
fi
while read -r j f; do
	[ "$j" = "insmod" ] || continue
	bn=${f##*/}
	un=${bn%.?z}
	un=${bn%.zst}
	[ -e "$CPIO/$un" ] && continue
	[ -n "$VERBOSE" ] && echo >&2 "initrd: Adding module $f"
	cp -au "$f" -t "$CPIO" 2>/dev/null || { RET=2; continue; }
	case "$bn" in
		*.gz)  gzip -qd "$CPIO/$bn" ;;
		*.xz)  xz   -qd "$CPIO/$bn" ;;
		*.zst) zstd -qd --rm "$CPIO/$bn" ;;
	esac
	printf '%s\n' "$un" >&3
done < "$CPIO/modprobe.txt" \
	3> "$CPIO/modules.conf"
rm "$CPIO/modprobe.txt"
cp -au /usr/lib/vm-run/initrd-init -T "$CPIO/init" || RET=2
mkdir -p "$CPIO"/etc
cp -a --parents /etc/os-release -t "$CPIO" &&
ln -s "os-release" "$CPIO/etc/initrd-release" || RET=2
for f in $FILES; do
	ff=${f%:*}
	ft=${f#"$ff"}
	ft=${ft#:}
	if [ -z "$ff" ]; then
		echo "initrd: Error: Cannot add empty filename for '$f'." >&2
		exit 1
	fi
	if [ ! -f "$ff" ] && [[ ! "$ff" =~ / ]]; then
		# Not directly accessible and does not contain '/' - perhaps executable from PATH?
		fx=$(type -p "$ff") && ff=$fx
		unset fx
	fi
	if [ ! -f "$ff" ]; then
		echo "initrd: Error: Cannot add '$ff': File not found." >&2
		exit 1
	fi
	if [ -x "$ff" ]; then
		ldd "$ff" >/dev/null 2>&1 && echo "initrd: Warning: $ff isn't a statically linked executable." >&2
		ftype=executable
	else
		ftype=regular
	fi
	[[ "$ft" =~ /$ ]] && ft+=$(basename "$ff")
	[ -n "$ft" ] || ft=$ff
	mkdir -p "$(dirname "$CPIO/$ft")"
	if [ -n "$VERBOSE" ]; then
		ftmsg=
		[ "$ff" = "$ft" ] || ftmsg=" -> $ft"
		echo >&2 "initrd: Adding $ftype file $ff$ftmsg"
	fi
	cp -a "$ff" -T "$CPIO/$ft" || RET=2
done

# Precreate box tools symlinks.
unset BOXBIN BOXLIST
if [ -x "$CPIO/bin/busybox" ]; then
	BOXBIN=bin/busybox
	BOXLIST=--list-full
elif [ -x "$CPIO/bin/toybox" ]; then
	BOXBIN=bin/toybox
	BOXLIST=--long
fi
if [ -v BOXBIN ]; then
	( cd "$CPIO"
	  mkdir -p {,usr/}{bin,sbin}
	  "$BOXBIN" "$BOXLIST" | xargs -n1 "--max-procs=$(nproc)" \
		  ln -sr "$BOXBIN" 2>/dev/null ||: )
fi

(cd "$CPIO" && find . | cpio -o -H newc --quiet -R 0:0 | gzip) > "$INITRD"
rm -rf "$CPIO"
exit $RET
