#!/bin/bash
# Script to regenerate initramfs images, e.g. on systemd updates
# May also be called from other packages or scripts
# Why it was developed:
# https://bugzilla.rosalinux.ru/show_bug.cgi?id=9516
# License: GPLv3
# Upstream at https://abf.io/import/rosa-kernel-tools
# Authors:
# - Mikhail Novosyolov <m.novosyolov@rosalinux.ru>, 2019-2021

set -efu

# mv(1) first copies the file and then deletes the origin;
# when there is not enough disk space, copying may fail;
# so fallback to _safe_mv if mv fails.
tmp= # prevent failing on unbound variable in trap
_safe_mv(){
	if [ -z "$1" ] || [ -z "$2" ]; then
		echo "arg1 or arg2 of _safe_mv is empty"
		return 1
	fi
	if [ -d "$1" ]
		then cp="cp -r"; mktemp="mktemp -d"
		else cp="cp"; mktemp="mktemp"
	fi
	tmp="$(${mktemp})"
	# try to fallback to permamnent storage if creating in /tmp
	# failed for some reason
	# TODO: better fallback if copying failed...
	if [ -z "$tmp" ]; then
		tmp="$(${mktemp} -p /root)"
	fi
	${cp} "$1" "$tmp"
	# TODO: if both $1 and $2 are directories, $1 will be put inside $2, not replaced
	${cp} --backup=none "$tmp" "$2"
	rm -fr "$1"
	rm -fr "$tmp"
}

_cleanup(){
	while read -r line0
	do
		# if file has already been deleted in another loop (line2)
		if [ ! -f "$line0" ]; then continue; fi
		line=
		filename=
		line="$(echo "$line0" | awk -F '.___SDIGSH_' '{print $1}')"
		filename="$(echo "$line" | sed -e 's,^/boot/,,g')"
		if [ -z "$line" ] || [ -z "$filename" ]; then continue; fi
		# if file exists and its size is > 0
		if [ -s "$line" ]
			then
				rm -f "$line0"
			else
				# If for some reason multiple backups exist, take the last one,
				# but verify that it is not empty
				# TODO: somehow verify validity of initramfs image not just by seing
				# that is is larger than zero bytes in size...
				DELETE_OTHER=0
				while read -r line2
				do
					if [ "$DELETE_OTHER" = 1 ]; then
						rm -f "$line2"
						continue
					fi
					if [ -s "$line2" ]; then
						mv "$line2" "$line" || _safe_mv "$line2" "$line"
						DELETE_OTHER=1
						continue
					fi
					# If it was not a file with size>0, but is a file,
					# so is a file with size=0, delete it.
					# As this script can't create anything other than a file,
					# don't delete it if it is not a file.
					if [ -f "$line2" ]; then
						rm -f "$line2"
					fi
				done < <(find /boot -maxdepth 1 -name "${filename}*.___SDIGSH_*" | sort -u)
		fi
	done < <(find /boot -maxdepth 1 -type f -name '*.___SDIGSH_*' -print)

	if [ -n "$tmp" ]; then rm -fr "$tmp"; fi
}
trap "_cleanup" EXIT

if systemd-detect-virt --quiet --container && ! env | grep -q '^SKIP_GEN=' ; then
	SKIP_GEN=1
fi

if [ -f /etc/initramfs-regen.conf ]; then
	. /etc/initramfs-regen.conf
fi

DEBUG="${DEBUG:-0}"
MSG="${MSG:-1}"
SKIP_GEN="${SKIP_GEN:-0}"
MAX_KERNELS="${MAX_KERNELS:-5}"
GENERATOR="${GENERATOR:-dracut}"

if ! command -v "$GENERATOR" >/dev/null 2>/dev/null; then
	export PATH="/usr/sbin:/sbin:${PATH}"
	if ! command -v "$GENERATOR" >/dev/null 2>/dev/null; then
		echo "Error: $GENERATOR not found!"
		exit 1
	fi
fi

if [ "$DEBUG" != 0 ]; then set -x; fi

if [ "$SKIP_GEN" = 1 ] ; then
# Note: we do have to regen initrd in chroot
# because chroots are often used in rescue purposes.
# But there is no need to spend time for this in containers
	if [ "$MSG" = 0 ]; then exit 0; fi
	echo "Running in container; no need to (re)generate initramfs — it's not used in containers."
	echo "To force (re)generating initramfs run:"
	echo "# env SKIP_GEN=0 initramfs-regen"
	echo "and/or add SKIP_GEN=0 to /etc/initramfs-regen.conf to make it permanent."
	exit 0
fi

# inspired a bit by grub_file_is_not_garbage() frpm grub2/util/grub-mkconfig_lib.in
# rpm-sort is an utility grub2-rpm-sort added by patches from Fedora
list_of_img="$(find /boot -maxdepth 1 \( -name 'initrd-*' -or -name 'initrd.img-*' -or -name 'initramfs-*' -or -name 'initramfs.img-*' \) \! \( -name '*.new' -or -name '*.rpmsave' -or -name '*.rpmnew' -or -name '*.dpkg-*' -or -iname '*.readme*' -or -iname '*.sig' \) -print | rpm-sort | tac)"
# TODO: here we could try to find all kernels and generate initramfs for them
# even if currently their initramfs does not exist, e.g. was deleted accidently.
# But then we must handle possible situation when we run out of disk space
# in the middle of process (because, when creating new initramfs, the probability
# of increasing disk usage is much higher then when just regenerating existing ones)
# and somehow guarantee that we don't leave our system without a working initramfs image.
# This seems to be not a trivial task that is out of scope of this script.
num_of_img="$(echo "$list_of_img" | wc -l)"

if [ "$num_of_img" -gt "$MAX_KERNELS" ] && [ "$MSG" -gt 0 ]; then
	echo "Initramfs for more than $MAX_KERNELS kernels found. That's too many and will take too long."
	echo "Will regenerate initramfs only for $MAX_KERNELS latest kernels."
	echo "You may set MAX_KERNELS= in /etc/initramfs-regen.conf if you don't like this and rerun initramfs-regen."
fi

count=0
while read -r line
do
	if [ -z "$line" ]; then continue; fi
	if [ "$count" -gt "$MAX_KERNELS" ]; then break; fi
	# There is no good way to quickly check if file is really and initramfs image or not
	# This situation must not happen, so let's explicitly print a warning if it does happen
	if [ ! -f "$line" ]; then
		echo "File $line does not exist or is not a file, skipping it."
		continue
	fi
	kernel=
	kernel="$(echo "$line" | sed \
		-e 's,^/boot/initrd-,,' \
		-e 's,^/boot/initrd.img-,,' \
		-e 's,^/boot/initramfs-,,' \
		-e 's,^/boot/initramfs.img-,,' \
		-e 's,\.img$,,' \
		)"
	if [ -z "$kernel" ]; then continue; fi
	# Now we must some how handle situations when the same /boot partition
	# is used by multiple co-installed Linux systems
	# and avoid regenerating initramfs images that don't belong to our system.
	# This situation is actually not handled in e.g. %post scripts of kernel packages...
	# But let's try to handle it.
	# Assume that _all_ kernels have their modules in /lib/modules
	# /lib/modules/<deleted kernel>/ may have subdirectories 'build' and 'source'
	# TODO: maybe find a better way
	if [ ! -d "/lib/modules/${kernel}/kernel" ]; then
		echo "$line seems to belong to another OS or kernel $kernel is not installed, skipping ${line}."
		continue
	fi
	if [ "$(ls "/lib/modules/${kernel}/kernel" | wc -l)" -lt 1 ]; then
		echo "Kernel $kernel seems to be removed or not fully installed or damaged, skipping $line"
		continue
	fi
	# Now make a backup of the initramfs that will be regenerated.
	# It has multiple purposes:
	# 1) validate that some disk space is free (regenerated images may be bigger that the old one)
	# 1) rollback if e.g. we run out of disk space while generating
	# trap will delete possibly left junk.
	# Here we assume that $GENERATOR itself takes care that it first generates initramfs
	# and only then replaces the old file with the new one, keeping /boot/* consistent.
	random="$(head -c 100 /dev/urandom | base64 | tr -cd '[:alnum:]._-' | head -c 10)"
	if [ -z "$random" ] || [ ${#random} -lt 10 ]
		# UNIX date is in the beginning to enable alphanum sorting
		then random="$(date +%s)_norandom"
		else random="$(date +%s)_${random}"
	fi
	postfix=".___SDIGSH_${random}"
	if ! cp "${line}" "${line}${postfix}"; then
		echo "Some error occured, probably out of disk space, skipping $line"
		continue
	fi
	SUCCESS=0
	case "$GENERATOR" in
		dracut )
			if ( set -x; dracut --force "$line" "$kernel" );
				then SUCCESS=1
			fi
		;;
		update-initramfs )
			if ( set -x; update-initramfs -u -k "$kernel" );
				then SUCCESS=1
			fi
		;;
		* )
			echo "Unknown \$GENERATOR, set dracut or update-initramfs. Exiting."
			exit 1
		;;
	esac
	if [ "$SUCCESS" = 1 ]; then
		count=$((count+1))
	fi
	# -s means that file exists and is not is not zero-sized
	if [ -s "$line" ]
		then
			rm -f "${line}${postfix}"
		else
			echo "Rolling back ${line}!"
			mv "${line}${postfix}" "${line}" || _safe_mv "${line}${postfix}" "${line}"
	fi
	unset kernel
done < <(echo "$list_of_img")
