#! /bin/bash -eu

# Maintain kernel-version-specific symlinks in /lib/firmware based on
# configuration present in /usr/share/microcode_ctl/ucode_with_caveats.
#
# SPDX-License-Identifier: CC0-1.0

usage()
{
	echo "Usage: update_ucode [--action {add|remove|refresh}]" \
	     "[--kernel KERNELVER]* [--verbose] [--dry-run]" \
	     "[--trigger-dracut] [--cleanup intel_ucode caveats_ucode]" >&2
}

debug() { [ 0 = "$verbose" ] || echo "$*" >&2; }

MC_DIR=/usr/share/microcode_ctl
INTEL_UCODE_DIR=intel-ucode
DATA_DIR=/usr/share/microcode_ctl/ucode_with_caveats
FW_DIR=/lib/firmware
check_caveats=/usr/libexec/microcode_ctl/check_caveats

action=refresh
kernel=
verbose=0
verbose_opt=
dry_run=0
trigger_dracut=0
remove_cleanup=0
cleanup_intel=
cleanup_caveats=

while [ 1 -le "$#" ]; do
	case "$1" in
	-a|--action)
		shift
		action="$1"
		;;
	-k|--kernel)
		shift
		kernel="$kernel $1"
		;;
	-v|--verbose)
		verbose=1
		verbose_opt="-v"
		;;
	-n|--dry-run)
		dry_run=1
		;;
	-d|--trigger-dracut)
		trigger_dracut=1
		;;
	-c|--cleanup)
		remove_cleanup=1
		shift
		cleanup_intel="$1"
		shift
		cleanup_caveats="$1"
		;;
	*)
		echo "Unknown argument \"$1\"" >&2
		usage
		exit 1
	esac
	shift
done

cmd=
[ 0 -eq "$dry_run" ] || cmd=echo

case "$action" in
add|remove|refresh)
	# Scan all directories in FW_DIR and all existing kernels
	if [ -z "$kernel" ]; then
		debug "No kernel versions provided, scanning..."

		kvers=$(find /lib/modules/ -name '[2-9].*' -print)
		for k_dir in $kvers; do
			k="${k_dir#/lib/modules/}"
			[ ! -e "/boot/symvers-$k.gz" ] || {
				debug "  Adding $k (from /lib/modules)"
				kernel="$kernel $k"
			}
		done

		kvers=$(find /lib/firmware/ -name '[2-9].*' -print)
		for k_dir in $kvers; do
			k="${k_dir#/lib/firmware/}"
			[ ! -d "$k_dir" ] || {
				debug "  Adding $k (from /lib/firmware)"
				kernel="$kernel $k"
			}
		done
	fi
	;;
*)
	echo "Unknown action \"$action\"" >&2
	usage
	exit 1
	;;
esac

# Generic part: managing intel ucode
debug "Running action \"$action\" on common Intel microcode directory"
while :; do
	[ ! -e "/etc/microcode_ctl/intel-ucode-disallow" ] || {
		debug "  Skipping \"$i\":" \
		      "\"/etc/microcode_ctl/intel-ucode-disallow\"" \
		      "present"
		break
	}
	[ ! -e "$FW_DIR/intel-ucode-disallow" ] || {
		debug "  Found \"$FW_DIR/intel-ucode-disallow\"," \
		      "skipping"
		break
	}

	# Removing old files
	case "$action" in
	refresh|remove)
		debug "  Removing old files from ${MC_DIR}/${INTEL_UCODE_DIR}"
		if [ 0 = "$remove_cleanup" ]; then
			find "${MC_DIR}/${INTEL_UCODE_DIR}" \
				-maxdepth 1 -mindepth 1 \
				-type f -printf '%f\n'
		else
			cat "$cleanup_intel"
		fi | while read -r fname; do
			name="${FW_DIR}/${INTEL_UCODE_DIR}/${fname}"

			# Needed in case we downgrade to a version where
			# no symlinks in /lib/firmware were used
			if [ 1 = "$remove_cleanup" ]; then
				[ -L "$name" ] || continue
			fi

			$cmd rm -f $verbose_opt "$name"
		done
		$cmd rmdir -p $verbose_opt "${FW_DIR}/${INTEL_UCODE_DIR}" 2>/dev/null || true
		;;
	esac

	# Adding new ones
	case "$action" in
	add|refresh)
		if grep -q '^flags[[:space:]]*:.*hypervisor' /proc/cpuinfo; then
			debug "  A virtualised environment has been detected" \
			      "(hypervisor flag is present in /proc/cpuinfo)," \
			      "skipping"
			break
		fi

		debug "  Creating symlinks in ${MC_DIR}/${INTEL_UCODE_DIR}"
		$cmd mkdir -p $verbose_opt "${FW_DIR}/${INTEL_UCODE_DIR}"
		$cmd find "${MC_DIR}/${INTEL_UCODE_DIR}" -maxdepth 1 -mindepth 1 \
			-type f -exec bash -c 'ln -s '"$verbose_opt"' '\''{}'\'' \
				"'"${FW_DIR}/${INTEL_UCODE_DIR}/"'$(basename '\''{}'\'')"' \;
		;;
	esac

	break
done

debug "Running action \"$action\" on kernels $kernel"

if [ 0 = "$remove_cleanup" ]; then
	ls "$DATA_DIR"
else
	cat "$cleanup_caveats"
fi | while read -r i; do
	debug "Processing data directory \"$i\"..."
	[ -e "$DATA_DIR/$i/readme" ] || {
		debug "  Skipping \"$i\": no readme"
		continue
	}
	[ -e "$DATA_DIR/$i/config" ] || {
		debug "  Skipping \"$i\": no config"
		continue
	}
	[ ! -e "/etc/microcode_ctl/ucode_with_caveats/disallow-$i" ] || {
		debug "  Skipping \"$i\":" \
		      "\"/etc/microcode_ctl/ucode_with_caveats/disallow-$i\"" \
		      "present"
		continue
	}

	for k in $(echo "$kernel"); do
		debug "    Processing kernel version \"$k\""
		$check_caveats -k "$k" -c "$i" $verbose_opt > /dev/null || {
			debug "    Checking for caveats failed" \
			      "(kernel version \"$k\"), skipping"
			continue
		}
		path=$($check_caveats -k "$k" -c "$i")

		case "$action" in
		remove|refresh)
			debug "    Removing $path (part of $action)..."

			if [ -e "$FW_DIR/$k/readme-$i" ]; then
				debug "      Removing \"$FW_DIR/$k/$path\""
				$cmd rm -f $verbose_opt "$FW_DIR/$k/$path"
				$cmd rm -f $verbose_opt "$FW_DIR/$k/readme-$i"
				$cmd rmdir -p $verbose_opt "$(dirname "$FW_DIR/$k/$path")" 2>/dev/null || true

				[ 0 -eq "$trigger_dracut" ] || {
					debug "      Triggering dracut for kernel \"$k\""
					$cmd dracut -f $verbose_opt --kver "$k"
				}
			else
				debug "      \"$FW_DIR/$k/readme-$i\" is not found," \
				      "skipping \"$FW_DIR/$k/$path\" removal"
			fi
			;;
		esac

		case "$action" in
		add|refresh)
			debug "    Adding $path (part of $action)..."

			[ -e "/boot/symvers-$k.gz" ] || {
				debug "      \"/boot/symvers-$k.gz\"" \
				      "does not exist, skipping"
				continue
			}
			if grep -q '^flags[[:space:]]*:.*hypervisor' /proc/cpuinfo; then
				debug "      A virtualised environment has been detected" \
				      "(hypervisor flag is present in /proc/cpuinfo)," \
				      "skipping"
				break
			fi
			[ ! -e "$FW_DIR/$k/disallow-$i" ] || {
				debug "      Found \"$FW_DIR/$k/disallow-$i\"," \
				      "skipping"
				continue
			}
			[ ! -e "$FW_DIR/$k/$path" ] || {
				debug "      \"$FW_DIR/$k/$path\" already exists," \
				      "skipping"
				continue
			}

			debug "      Adding \"$FW_DIR/$k/$path\""
			$cmd mkdir -p "$(dirname "$FW_DIR/$k/$path")"
			$cmd ln -s "$DATA_DIR/$i/$path" "$FW_DIR/$k/$path"
			$cmd cp "$DATA_DIR/$i/readme" "$FW_DIR/$k/readme-$i"

			[ 0 -eq "$trigger_dracut" ] || {
				debug "      Triggering dracut for kernel \"$k\""
				$cmd dracut -f $verbose_opt --kver "$k"
			}
			;;
		remove)
		esac
	done
done