#!/bin/sh
# busybox based initramfs script

add_kmod() {
	while [ "$1" ]; do
		kmod=$(echo "$1" | tr - _); shift
		kmodpath=$(modprobe -S "$K" -l "$kmod" 2>/dev/null) || continue
		path=/lib/modules/$K/${kmodpath}
		[ -f "$W/$path" ] && continue
		add_file "$path"
		for firmware in $(modinfo -k "$K" -F firmware "$kmod" 2>/dev/null); do
			# firmware may be a glob pattern (e.g. brcm/brcmbt4377*.bin),
			# so leave it unquoted here and quote each matched file
			for fw in /lib/firmware/$firmware.xz /lib/firmware/$firmware.gz /lib/firmware/$firmware; do
				[ -f "$fw" ] && add_file "$fw"
			done
		done
		for depends in $(modprobe -S "$K" -D "$kmod" 2>/dev/null | awk '{print $2}'); do
			[ "$depends" ] || continue
			depends=${depends##*/}; depends=${depends%.ko*}
			add_kmod "$depends"
		done
	done
}

# returns non-zero if any binary is missing
add_bin() {
	ret=0
	while [ "$1" ]; do
		# search the original PATH, not busybox's applet symlinks
		bin=$(PATH=${P:-$PATH} which "$1") || {
			echo "FAILED: $1: missing binary"
			ret=1
			shift
			continue
		}
		shift
		add_file "$bin"
		[ -L "$bin" ] && bin=$(readlink -f "$bin")
		# shared libraries and the dynamic loader, if any
		for file in $(ldd "$bin" 2>/dev/null | grep -o '/[^ ]*'); do
			add_file "$file"
		done
	done
	return $ret
}

add_file() {
	f=$1
	t=$2
	m=$3
	# skip already copied file
	[ -e "$W${t:-$f}" ] && return
	[ -f "$f" ] || {
		echo "FAILED: $f: missing file"
		return
	}
	if [ -L "$f" ]; then
		mkdir -p "${W}${f%/*}"
		ln -sf "$(readlink "$f")" "${W}${f}"
		f=$(readlink -f "$f")
	fi
	install ${m:+-m "$m"} -D "$f" "$W${t:-$f}" 2>/dev/null
	echo "$([ "$?" = 0 ] && echo OK || echo FAILED): $f"
}

# print hooks in a directory, without extension
list_hooks() {
	for hook in "$1"/*.sh "$1"/*.pre "$1"/*.post; do
		[ -f "$hook" ] && echo "${hook%.*}"
	done | sort -u
}

# source <hook>.sh at build time, then add <hook>.pre and <hook>.post to the
# image to run at boot. A .sh hook sets SKIP=1 to leave those out.
run_hook() {
	hook=$1
	SKIP=
	if [ -f "$hook.sh" ]; then
		echo "RUNNING HOOK: $hook.sh..."
		. "$hook.sh"
	fi
	[ "$SKIP" ] && return
	echo "${hook##*/}" >> "$W"/hooks/order
	for stage in pre post; do
		[ -f "$hook.$stage" ] && add_file "$hook.$stage" "/hooks/${hook##*/}.$stage"
	done
}

K=$(uname -r)
I=/usr/share/mkinitrd/init
HD=/usr/share/mkinitrd/hooks.d
HS=/usr/share/mkinitrd/hooks
C=/etc/mkinitrd.conf
CD=/etc/mkinitrd.d

while [ "$1" ]; do
	case $1 in
		-k|-i|-c)
			[ "$2" ] || { echo "error: option '$1' needs an argument"; exit 1; }
			case $1 in -k) K=$2;; -i) I=$2;; -c) C=$2; CC=1;; esac
			shift;;
		-f) F=1;;
		-h) H=1;;
		-*) echo "error: unknown option '$1'"; exit 1;;
		 *) O=$1;;
	esac
	shift
done

if [ "$H" ]; then
	cat <<EOF
mkinitrd: busybox based initramfs generator.

Usage: $0 [-fh] [-c config] [-k kernel] [-i init] [ OUTPUT ]

	-c config   Use custom config file (default: /etc/mkinitrd.conf)
	-f          Replace existing output file
	-i init     Use custom init file (default: /usr/share/mkinitrd/init)
	-k kernel   kernel version
	-h          Show this help message
EOF
exit 0
fi

if [ "$(id -u)" != 0 ]; then
	echo "error: this program need to run as root"
	exit 1
fi

if [ -f "$C" ]; then
	. "$C"
elif [ "$CC" ]; then
	echo "error: config file '$C' not found"
	exit 1
fi

# OUTPUT argument takes precedence over INITRD in config
O=${O:-${INITRD:-initrd}}

# replace the file a symlink points to, not the symlink
[ -L "$O" ] && O=$(readlink -f "$O")

if [ ! -d "$(dirname "$O")" ]; then
	echo "error: output directory '$(dirname "$O")' not found"
	exit 1
fi

# the image is renamed into place, which would replace a device or directory
if [ -e "$O" ] && [ ! -f "$O" ]; then
	echo "error: '$O' is not a regular file"
	exit 1
fi

if [ -e "$O" ] && [ ! "$F" ]; then
	echo "error: '$O' already exists, use -f to replace it"
	exit 1
fi

# hooks selected in config, mdev if HOOKS is not set at all
HOOKS=${HOOKS-mdev}

for h in $HOOKS; do
	if [ ! -f "$HS/$h.sh" ] && [ ! -f "$HS/$h.pre" ] && [ ! -f "$HS/$h.post" ]; then
		echo "error: hook '$h' not found in $HS"
		exit 1
	fi
done

M=/lib/modules/$K

if [ ! -d "$M" ]; then
	echo "error: kernel modules directory '$M' not found"
	exit 1
fi

if [ ! -f "$M"/modules.dep ]; then
	echo "error: modules.dep file not found, please run 'depmod' first"
	exit 1
fi

if [ ! "$(command -v busybox)" ]; then
	echo "error: busybox not found in PATH, please install busybox"
	exit 1
fi

if [ ! -f "$I" ]; then
	echo "error: init file '$I' not found"
	exit 1
fi

# work dir and temporary bin path for busybox's applet
# we trying to use busybox utils as possible
W=$(mktemp -d) && B=$(mktemp -d) || exit 1
T=$O.tmp.$$
trap 'rm -rf "$W" "$B" "$T"' EXIT
trap 'exit 1' INT TERM HUP
chmod 755 "$W"
for a in $(busybox --list); do
	ln -s "$(command -v busybox)" "$B/$a"
done
P=$PATH
export PATH="$B:$PATH"

mkdir -p "$W/$M"

for d in bin dev etc hooks lib proc run sys; do
	mkdir -p "$W/$d"
done
ln -s . "$W"/usr
ln -s bin "$W"/sbin
ln -s lib "$W"/lib64

# modules loaded on this host, or a minimal set to boot common hardware
case $HOSTONLY in
	[Yy][Ee][Ss]|1|[Tt][Rr][Uu][Ee])
		for i in /sys/module/*; do
			add_kmod "${i##*/}"
		done
		;;
	*)
		add_kmod \
			virtio_pci virtio_blk virtio_scsi nvme ahci ata_piix \
			sd_mod sr_mod mmc_block sdhci_pci \
			xhci_pci ehci_pci ohci_pci uhci_hcd usb_storage uas \
			usbhid hid_generic atkbd i8042 \
			ext4 xfs btrfs f2fs vfat
		;;
esac

# extra modules from config
[ "$MODULES" ] && add_kmod $MODULES

# root filesystem from config, init uses it instead of auto
if [ "$ROOTFSTYPE" ]; then
	add_kmod "$ROOTFSTYPE"
	echo "$ROOTFSTYPE" > "$W"/etc/rootfstype
fi
add_file "$M"/modules.builtin
add_file "$M"/modules.order

add_bin busybox

# extra binaries and files from config
[ "$BINARIES" ] && add_bin $BINARIES
for f in $FILES; do
	case $f in
		*:*) add_file "${f%%:*}" "${f#*:}";;
		  *) add_file "$f";;
	esac
done
install -m755 "$I" "$W"/init

for h in $(list_hooks "$HD") $(for h in $HOOKS; do echo "$HS/$h"; done) $(list_hooks "$CD"); do
	run_hook "$h"
done

# after hooks, so modules added by them are indexed too
depmod -b "$W" "$K"

# write next to the output and rename, so a failure keeps the existing image
if ! ( cd "$W" ; find . | cpio -o -H newc --quiet | gzip -9 ) > "$T"; then
	echo "error: failed to write '$O'"
	exit 1
fi
[ -f "$O" ] && chmod "$(stat -c %a "$O")" "$T"
if ! mv -f "$T" "$O"; then
	echo "error: failed to write '$O'"
	exit 1
fi
du -sh "$O"

exit 0
