#!/bin/sh
#
# a simple initramfs generator inspired by Arch's mkinitcpio
# but wrote it simpler
#

PATH=/bin:/usr/bin:/sbin:/usr/sbin

# msg/msgerr/msgwarn:
#   - msgerr: always writes to stderr, NEVER silenced (a critical error must always be visible)
#   - msgwarn: respects -q so "mkinitramfs -q" suppresses non-fatal warnings like missing firmware
#   - msg: informational, respects -q
msg() {
	[ "$QUIET" ] && return
	printf ':: %s\n' "$*"
}

msgerr() {
	printf 'ERROR: %s\n' "$*" >&2
}

msgwarn() {
	[ "$QUIET" ] && return
	printf 'WARNING: %s\n' "$*" >&2
}

add_binary() {
	binpath=$(command -v "$1")
	if [ -z "$binpath" ]; then
		msgwarn "missing binary: $1"
		return
	fi
	add_file "$binpath"
	if [ -L "$binpath" ]; then
		add_binary "$(readlink -f "$binpath")"
		return
	fi
	case "$(file -bi "$binpath")" in
		*application/x-sharedlib* | *application/x-executable* | *application/x-pie-executable*)
			binlib=$(ldd "$binpath" | awk 'NF==4 && $3 ~ /^\// { print $3 }')
			for lib in $binlib; do
				add_file "$lib"
			done;;
	esac
	unset binpath binlib
}

add_module() {
	if modinfo -k "$KERNEL" "$1" >/dev/null 2>&1; then
		modname=$(modinfo -k "$KERNEL" -F name     "$1" | head -n1)
		modpath=$(modinfo -k "$KERNEL" -F filename "$1" | head -n1)
		[ -z "$modname" ] && return 0
		[ -z "$modpath" ] && return 0
	else
		msg "missing module: $1"
		return
	fi
	if [ -f "$INITDIR/lib/modules/$KERNEL/kernel/${modpath##*/}" ]; then
		return
	fi
	add_file "$modpath" "lib/modules/$KERNEL/kernel/${modpath##*/}"
	modinfo -F firmware -k "$KERNEL" "$modname" | while read -r line; do
		if [ ! -f "/lib/firmware/$line" ]; then
			msgwarn "missing firmware for $modname: $line"
		else
			add_file "/lib/firmware/$line"
		fi
	done
	for i in $(modinfo -F depends -k "$KERNEL" "$modname" | tr ',' ' '); do
		add_module "$i"
	done
	unset modname modpath
}

add_file() {
	[ "$1" ] || return

	src=$1

	case $src in
		/*) ;;
		*)  msgwarn "absolute source path needed: $src"; return ;;
	esac

	if [ ! -f "$src" ]; then
		msgwarn "file not found: $src"
		return
	fi

	if [ -z "$2" ]; then
		dest="${src#/}"
	else
		dest="$2"
		# Now the leading-slash guard operates on the value actually supplied
		case $dest in
			/*) msgwarn "destination path must without leading '/': $dest"
				unset src dest
				return ;;
		esac
	fi

	mode=${3:-$(stat -c %a "$src")}
	if [ -z "$mode" ]; then
		msgwarn "failed get file mode: $src"
		unset src dest
		return
	fi

	install -Dm"$mode" "$src" "$INITDIR/$dest"

	unset src dest mode
}

add_dir() {
	path=$1
	mode=${2:-755}

	if [ -z "$path" ]; then
		return 1
	fi

	case $path in
		/*) ;;
		 *) return 1;;
	esac

	if [ -d "$INITDIR$path" ]; then
		return 0
	fi

	install -dm"$mode" "$INITDIR$path"

	unset path
}

add_symlink() {
	name=$1
	target=$2

	[ "$name" ] || return 1

	if [ ! "$target" ]; then
		target=$(readlink -f "$name")
		if [ ! "$target" ]; then
			msgerr "invalid symlink: $name"
			return 1
		fi
	fi

	add_dir "${name%/*}"

	if [ -L "$INITDIR$name" ]; then
		msgwarn "overwriting symlink: $name"
	fi

	ln -sfn "$target" "$INITDIR$name"
}

finalize_modules() {
	[ -d "$INITDIR/lib/modules/$KERNEL/kernel" ] || return
	for file in /lib/modules/"$KERNEL"/modules.*; do
		add_file "$file"
	done
	awk -F'/' '{ print "kernel/" $NF }' "/lib/modules/$KERNEL/modules.order" \
		> "$INITDIR/lib/modules/$KERNEL/modules.order"
	depmod -b "$INITDIR" "$KERNEL"
}

get_hook_path() {
	case $1 in
		*/*) [ -f "$1" ] && printf '%s\n' "$1" ;;
		*)   for i in $(printf '%s\n' "$HOOKDIRS" | tr ':' ' '); do
				if [ -f "$i/$1" ]; then
					printf '%s\n' "$i/$1"
					break
				fi
			 done ;;
	esac
}

run_build_hook() {
	if [ "$HOOKS" ]; then
		for hook in $HOOKS; do
			case " $DONEHOOK " in
				*" $hook "*) continue ;;
			esac
			hookpath=$(get_hook_path "$hook.hook")
			if [ "$hookpath" ]; then
				. "$hookpath"
				if [ "$(command -v build_hook)" ]; then
					msg "running build_hook: $hook"
					build_hook
				fi
				printf '%s\n' "$hook" >> "$INITDIR/hook/hook.order"
				add_file "$hookpath" "hook/$hook" 755
				DONEHOOK="$DONEHOOK $hook"
				unset hookpath build_hook
			else
				msgwarn "missing hook: $hook"
			fi
		done
	fi
}

cmd_exists() {
	command -v "$1" >/dev/null 2>&1 || msgerr "command '$1' is not available."
}

add_compression() {
	complevel_gzip="${complevel_gzip:--9}"
	complevel_xz="${complevel_xz:--9}"
	complevel_zstd="${complevel_zstd:--19}"
	complevel_lz4="${complevel_lz4:--12}"

	case "$INITFSCOMP" in
		gzip) comp="$(command -v gzip 2>/dev/null || printf 'gzip') $complevel_gzip" ;;
		xz)   cmd_exists xz;   comp="xz -C crc32 -T0 $complevel_xz" ;;
		zstd) cmd_exists zstd; comp="zstd -T0 $complevel_zstd" ;;
		lz4)  cmd_exists lz4;  comp="lz4 --favor-decSpeed -lz $complevel_lz4" ;;
		none) comp="cat" ;;
		*)    comp="gzip $complevel_gzip" ;;
	esac
}

cleanup() {
	rm -fr "$INITDIR"
}

interrupted() {
	cleanup
	exit 1
}

usage() {
	cat << EOF
Usage:
  $(basename "$0") [option] [argument]

Options:
  -k <version>     custom kernel version (default: $KERNEL)
  -o <output>      custom output name (default: $INITRAMFS)
  -i <init>        custom init file (default: $INITIN)
  -m <modules>     add extra modules (comma separated)
  -b <binaries>    add extra binary (comma separated)
  -f <file>        add extra file (comma separated & absolute path)
  -c <config>      use custom config (default: $CONFIG)
  -C <compression> (gzip|xz|zstd|lz4|none  default: $INITFSCOMP)
  -A <hook>        add extra hook (comma separated, precedence over -a, -s & HOOKS)
  -a <hook>        add extra hook (comma separated, precedence over -s & after HOOKS)
  -s <hook>        skip hook defined in HOOKS (comma separated)
  -q               quiet mode
  -h               print this help msg

EOF
}

needarg() {
	if [ ! "$1" ]; then
		printf 'ERROR: argument is needed for this option!\n' >&2
		exit 1
	fi
}

parse_opt() {
	while [ "$1" ]; do
		case $1 in
			-k) needarg "$2"; KERNEL=$2; shift ;;
			-o) needarg "$2"; OUTPUT=$2; shift ;;
			-i) needarg "$2"; INITIN=$2; shift ;;
			-c) needarg "$2"; CONFIG=$2; shift ;;
			-C) needarg "$2"; INITFSCOMP=$2; shift ;;
			-A) needarg "$2"; ADDEARLYHOOKS=$(printf '%s\n' "$2" | tr ',' ' '); shift ;;
			-a) needarg "$2"; ADDHOOKS=$(printf '%s\n' "$2"      | tr ',' ' '); shift ;;
			-s) needarg "$2"; SKIPHOOKS=$(printf '%s\n' "$2"     | tr ',' ' '); shift ;;
			-m) needarg "$2"; ADDMODULES=$(printf '%s\n' "$2"    | tr ',' ' '); shift ;;
			-b) needarg "$2"; ADDBINARIES=$(printf '%s\n' "$2"   | tr ',' ' '); shift ;;
			-f) needarg "$2"; ADDFILES=$(printf '%s\n' "$2"      | tr ',' ' '); shift ;;
			-q) QUIET=1 ;;
			-h) usage; exit 0 ;;
			*)  printf 'ERROR: invalid option %s\n' "'$1'" >&2; exit 1 ;;
		esac
		shift
	done
}

main() {
	parse_opt "$@"

	if [ "$(id -u)" != "0" ]; then
		msgerr "need root access!"
		exit 1
	fi

	CPIO=$(command -v bsdcpio) || CPIO=$(command -v cpio)
	if [ ! "$CPIO" ]; then
		msgerr "neither 'bsdcpio' nor 'cpio' found, exiting..."
		exit 1
	fi

	if [ -f "$CONFIG" ]; then
		. "$CONFIG"
	else
		msgerr "config file '$CONFIG' does not exist."
		exit 1
	fi

	if [ ! -d "/lib/modules/$KERNEL" ]; then
		msgerr "kernel directory '/lib/modules/$KERNEL' does not exist."
		exit 1
	fi

	if [ ! -f "$INITIN" ]; then
		msgerr "init file '$INITIN' does not exist."
		exit 1
	fi

	if [ "$OUTPUT" ]; then
		if [ "$(basename "$OUTPUT")" != "$OUTPUT" ] && [ ! -d "$(dirname "$OUTPUT")" ]; then
			msgerr "directory '$(dirname "$OUTPUT")' for output '$(basename "$OUTPUT")' does not exist."
			exit 1
		elif [ -d "$OUTPUT" ]; then
			msgerr "'$OUTPUT' is a directory."
			exit 1
		fi
		INITRAMFS="$OUTPUT"
	fi

	# filter out skip hooks (-s)
	if [ "$SKIPHOOKS" ] && [ "$HOOKS" ]; then
		for h in $HOOKS; do
			printf '%s\n' "$SKIPHOOKS" | tr ' ' '\n' | grep -qx "$h" || NEWHOOKS="$NEWHOOKS $h"
		done
		HOOKS=$NEWHOOKS
		unset NEWHOOKS
	fi

	# add extra hooks (-a)
	[ "$ADDHOOKS" ]      && HOOKS="$HOOKS $ADDHOOKS"

	# add extra early hooks (-A)
	[ "$ADDEARLYHOOKS" ] && HOOKS="$ADDEARLYHOOKS $HOOKS"

	# add extra modules/files/binaries
	[ "$ADDMODULES"   ] && MODULES="$MODULES $ADDMODULES"
	[ "$ADDFILES"     ] && FILES="$FILES $ADDFILES"
	[ "$ADDBINARIES"  ] && BINARIES="$BINARIES $ADDBINARIES"

	[ "$QUIET" ] || printf 'Generating initramfs...\n'

	mkdir -p "$INITDIR/hook" "$INITDIR/newroot"
	install -m0755 "$INITIN" "$INITDIR/init"
	run_build_hook

	if [ "$BINARIES" ]; then
		msg "adding extra binaries..."
		for b in $BINARIES; do
			add_binary "$b"
		done
	fi

	if [ "$MODULES" ]; then
		msg "adding extra modules..."
		for m in $MODULES; do
			add_module "$m"
		done
	fi

	if [ "$FILES" ]; then
		msg "adding extra files..."
		for f in $FILES; do
			add_file "$f"
		done
	fi

	add_compression
	finalize_modules

	msg "generating initramfs image..."
	rm -f "$INITRAMFS"
	( cd "$INITDIR" || exit 1; find . | LANG=C "$CPIO" -o -H newc --quiet | $comp ) > "$INITRAMFS"

	cleanup
	[ "$QUIET" ] || printf 'Done: %s (%s)\n' "$INITRAMFS" "$(du -h "$INITRAMFS" | awk '{print $1}')"

	exit 0
}

trap interrupted 1 2 3 15

INITDIR="/tmp/mkinitramfs.$$"
KERNEL="$(uname -r)"
INITIN="/usr/share/mkinitramfs/init.in"
INITRAMFS="initrd-$KERNEL.img"
CONFIG="/etc/mkinitramfs.conf"
HOOKDIRS="/usr/share/mkinitramfs/hooks"

main "$@"
