#!/bin/ksh93

typeset -r VERSION='1.0' FPROG=${.sh.file} PROG=${FPROG##*/} SDIR=${FPROG%/*} \
	EFI_PATH='/boot/efi' KERNEL_PATH='/boot'

typeset -T LogObj_t=(
	typeset -Sh 'Color for info messages' GREEN='38;5;232;48;5;118'	#'1;30;102';
	typeset -Sh 'Color for warning messages' BLUE='38;5;21;48;5;118' #'1;34;102';
	typeset -Sh 'Color for fatal messages' RED='38;5;9;48;5;118' #'1;31;102';
	function log {
		print -u2 "\E[1;$2m${ date +%T; } $1:\E[0m $3"
	}
	typeset -Sfh ' log a message' log
	function info {
		_.log "INFO" ${_.GREEN} "$*"
	}
	typeset -Sfh ' log a info message' info
	function warn {
		_.log "WARN" ${_.BLUE} "$*"
	}
	typeset -Sfh ' log a warning message' warn
	function fatal {
		_.log "FATAL" ${_.RED} "$*"
	}
	typeset -Sfh ' log a fatal error message' fatal
	function printMarker {
		typeset COLOR="$1"
		print -f '\E[1;%sm----------------------------------------------------------------------------\E[0m\n' "${COLOR:-${_.GREEN}}"
	}
	typeset -Sfh ' print a marker line' printMarker
)
LogObj_t Log

function showUsage {
	[[ -n $1 ]] && X='-?' ||  X='--man'
	getopts -a ${PROG} "${ print ${USAGE} ; }" OPT $X
}

function checkMachineID {
	if [[ -f /etc/machine-id ]]; then
		OPT['MID']=$(</etc/machine-id)
		[[ ${OPT['MID']} =~ ^[0-9a-f]+$ ]] && return 0
		Log.fatal "Invalid imachine ID '${OPT[MID]}'. Exiting."
		return 2
	fi
	Log.fatal '/etc/machine-id not found.' \
		'\n\tUse systemd-machine-id-setup or something else to create one.' \
		'\n\tExiting.'
	return 1
}

function checkEFI {
	[[ -z ${OPT['EFI']} ]] && OPT['EFI']=${EFI_PATH}
	typeset EFI="${OPT['EFI']}"

	[[ -d ${EFI} ]] || { Log.fatal "EFI path '${EFI}' not found." \
		'Exiting.' ; return 2 ; }
	[[ -e ${EFI}/EFI/systemd/systemd-bootx64.efi ]] && return 0
	Log.fatal "'${EFI}/EFI/systemd/systemd-bootx64.efi' not found." \
		"\n\tDid you already run 'bootctl install'?"
	return 3
}

function genConfig {
	typeset R="$1" FMT="%-10s %s\n"
	printf "${FMT}" 'title'			"$2"
	printf "${FMT}" 'version'		"$R"
	printf "${FMT}" 'machine-id'	${OPT['MID']}
	printf "${FMT}" 'options'		"${OPT['BOPT']}"
	printf "${FMT}" 'linux'			"/${OPT['MID']}/$R/linux"
	printf "${FMT}" 'initrd'		"/${OPT['MID']}/$R/initrd"
}

function doAdd {
	typeset -A ALL TBD
	typeset MID="${OPT['MID']}" EFI="${OPT['EFI']}" KPATH="${OPT['KERNEL']}" COP
	typeset F X R T DRY= TITLE=${ lsb_release -ds ; } BOPT="${OPT['BOPT']}"
	integer VERB=0 DEPMOD=0 ERR=0
	[[ -n ${OPT['VERB']} ]] && VERB=1
	[[ -n ${OPT['DRY']} ]] && DRY='print -- '
	[[ -n ${OPT['DEPMODE']} ]] && DEPMOD=1
	[[ -z ${KPATH} ]] && KPATH=${KERNEL_PATH}

	for X in ${EFI}/{${MID},loader,loader/entries} ; do
		[[ -d $X ]] && continue
		if ! ${DRY} mkdir "$X" ; then
			Log.fatal "Unabe to create '$X'. Exiting."
			return 1
		fi
	done
	for R in ${OPT['ADD']} ; do
		[[ -z $R ]] && continue
		if [[ ! $R =~ ^[-._0-9a-zA-Z]+$ ]]; then
			Log.warn "Invalid release name '$R' ignored."
			continue
		fi
		[[ $R == 'current' ]] && R=${ uname -r; }
		TBD["_$R"]=1 
	done
	if [[ -n ${TBD['_all']} ]]; then
		unset TBD; typeset -A TBD
		for F in ~(N)${KPATH}/initrd.img-* ; do
			X=${F##*/}
			[[ ${X:11} =~ ^[-._0-9a-zA-Z]+$ ]] || continue
			[[ -f ${F%/*}/vmlinuz-${X:11} ]] || continue
			TBD["_${X:11}"]=1
		done
	fi
	[[ -e /proc/cmdline ]] && { X=$(</proc/cmdline) && COPT="${X#* }"; } || \
		COPT='ro root=ZFS=rpool/ROOT/linux'	# fallback
	[[ -z ${OPT['TRY']} ]] && OPT['TRY']=1
	# We try to change not more than really is needed.
	for R in ${!TBD[@]} ; do
		R=${R:1}
		F="${KPATH}/initrd.img-$R"
		if [[ ! -e $F ]]; then
			Log.warn "'$F' not found - skipping '$R'."
			continue
		fi
		F="${KPATH}/vmlinuz-$R"
		if [[ ! -e $F ]]; then
			Log.warn "'$F' not found - skipping '$R'."
			continue
		fi
		if (( DEPMOD )) && [[ -d /lib/modules/$R/kernel ]]; then
			(( VERB )) && Log.info "Running 'depmod -a $R' ..."
			${DRY} depmod -a "$R"
		fi
		F="${EFI}/${MID}/$R"
		if [[ ! -d $F ]] && !  ${DRY} mkdir $F ; then
			Log.fatal "Unable to create '$F' - skipping '$R'."
			(( ERR++ ))
			continue
		fi
		(( VERB )) && Log.info "Update entry for '$R' ..."
		F="${EFI}/${MID}/$R/initrd"
		[[ -e $F ]] && cmp -s $F ${KPATH}/initrd.img-$R
		if (( $? )) && ! ${DRY} cp -p ${KPATH}/initrd.img-$R $F ; then
			Log.warn "Copy ${KPATH}/initrd.img-$R to $F failed - skipping '$R'."
			(( ERR++ ))
		fi
		F="${EFI}/${MID}/$R/linux"
		[[ -e $F ]] && cmp -s $F ${KPATH}/vmlinuz-$R
		if (( $? )) && ! ${DRY} cp -p ${KPATH}/vmlinuz-$R $F ; then
			Log.warn "Copy ${KPATH}/vmlinuz-$R to $F failed - skipping '$R'."
			(( ERR++ ))
		fi
		F="${EFI}/loader/entries/${MID}-${R}"
		(( OPT['TRY'] > 1 )) && F+="+${OPT['TRY']}"
		F+=".conf"
		if [[ -z ${BOPT} ]]; then
			if [[ -e $F ]]; then
				# keep options as is
				while read X T ; do
					[[ $X == 'options' ]] && OPT['BOPT']="$T" && break
				done<$F
			else
				# use the ones from this kernel boot
				OPT['BOPT']="${COPT}"
			fi
		fi
		X=${ genConfig "$R" "${TITLE:-Linux} - $R"; }
		if [[ -e $F ]]; then
			T=$(<$F)
			if [[ $X == $T ]]; then
				(( VERB )) && Log.info 'done.'
				continue
			fi
		fi
		if (( OPT['TRY' ] )); then
			[[ -e ${F%+*}.conf ]] && ${DRY} rm -f ${F%+*}.conf
			for X in ~(N)${F%+*}+*.conf ; do
				[[ $X == $F ]] && continue
				${DRY} rm -f "$X" || Log.warn "Unable to remove '$X'."
			done
		fi
		if [[ -n ${DRY} ]]; then
			${DRY} print "$X" ' > ' $F
		else
			print "$X" >$F
			if (( $? )); then
				(( ERR++ ))
				Log.warn "Failed to write '$F'." && continue
			fi
		fi
		(( VERB )) && Log.info 'done.'
	done
	return ${ERR}
}

function doDelete {
	typeset -A ALL TBD
	typeset X R MID="${OPT['MID']}" EFI="${OPT['EFI']}" DRY=
	integer VERB=0
	[[ -n ${OPT['VERB']} ]] && VERB=1

	for R in ~(N)${EFI}/loader/entries/${MID}-*.conf ; do
		X=${R##*/}
		X=${X:${#MID}+1:${#X}-${#MID}-6}
		ALL[_${X%+*}]=1
	done
	for R in ~(N)${EFI}/${MID}/* ; do
		[[ -d $R ]] || continue
		X=${R##*/}
		ALL[_$X]=1
	done
	for R in ${OPT['DEL']} ; do
		[[ -z $R ]] && continue
		if [[ ! $R =~ ^[-._0-9a-zA-Z]+$ ]]; then
			Log.warn "Invalid release name '$R' ignored."
			continue
		fi
		[[ $R == 'current' ]] && R=${ uname -r; }
		TBD["_$R"]=1 
	done
	[[ -n ${TBD['_all']} ]] && typeset -n DO=ALL || typeset -n DO=TBD
	[[ -n ${OPT['DRY']} ]] && DRY='print -- '
	for R in ${!DO[@]} ; do
		(( VERB )) && Log.info "Removing '${R:1}' ..."
		if [[ -z ${ALL["$R"]} ]]; then
			(( VERB )) && Log.info "Release '$R' ignored - no entry found."
			continue
		fi
		${DRY} rm -f ${EFI}/loader/entries/${MID}-${R:1}.conf
		${DRY} rm -f ${EFI}/loader/entries/${MID}-${R:1}+[0-9]*.conf
		${DRY} rm -rf ${EFI}/${MID}/${R:1}
	done
}

function doMain {
	integer ERR=0

	[[ -z ${OPT['ADD']} && -z ${OPT['DEL']} ]] && { bootctl list ; return 0; }

	checkMachineID || (( ERR++ ))
	checkEFI || (( ERR++ ))
	(( ERR )) && return 1

	[[ -n ${OPT['DEL']} ]] && doDelete


	[[ -n ${OPT['ADD']} ]] && doAdd
}

USAGE="[-?${VERSION}"' ]
[-copyright?Copyright (c) 2021 Jens Elkner. All rights reserved.]
[-license?CDDL 1.0]
[+NAME?'"${PROG}"' - update EFI boot environment wrt. available Linux kernels.]
[+DESCRIPTION?This simple script checks the related EFI boot filesystm wrt. to the corresponding linux kernels and RAM disks and updates them accordingly. If neither \b-a\b nor \b-d\b is given, the current EFI boot env will be shown, only. If both are given, first \b-d\b and than \b-a\b option gets processed.]
[+?Before executing this script, one should run \bupdate-initramfs\b(8) to update the RAM disk image, which gets coppied over to the EFI filesystem by this script.]
[+?This scripts needs to be run by a user with read/write privileges for EFI filesystem (which must be already mounted) as well as read privileges for the kernel and RAM disk image, which needs to be copied to the EFI filesystem, so on most systems as user root.]
[h:help?Print this help and exit.]
[F:functions?Print a list of all functions available.]
[T:trace]:[functionList?A comma separated list of functions of this script to trace (convinience for troubleshooting).] 
[+?]
[D:depmod?Run "depmod -a" for the given release before add a new entry to the boot environment. This is usually not needed, for convenience, only.]
[a:add]:[release?Add the kernel with the given \arelease\a to the EFI boot environment. Can be given multiple times. Two special words are allowed as well: \bcurrent\b gets replaced with the release of the currently running kernel  and \ball\b gets replaced by all kernel releases found in the related directory.]
[d:delete]:[release?Delete the kernel with the \arelease\a from the EFI boot environment. Can be given multiple times. The special words \bcurrent\b and \ball\b are allowed as well and have the same meaning as for \b-a ...\b.]
[e:efiboot]:[path?Assume the given \apath\a represents the EFI filesystem in use. No check is made, whether the related filesystem is actually vfat, nor whether it is consistant with /etc/fstab entries, etc.. Default: \b'"${EFI_PATH}"'\b]
[k:kernels]:[path?Lookup intended kernels (\blinux-*\b) and related RAM disks (\binitrd.img-*\b) in the given path. Default: \b'"${KERNEL_PATH}"'\b]
[n:dry?Just show, what would be done but do not really do it.]
[o:option]:[option?Pass the given \aoption\a to the kernel to start. Can be used multiple times and as quoted string at once. If none is given, those from the boot configuration for the related kernel will be used. If there is no configuration yet, the options used to start the currently running kernel will be used.]
[t:tries]:[num?Try \anum\a times to start the kernel. Non-Integer values get silently ignored. Default: 1]
[v:verbose?Emit the boring details about what the script is doing, too.]
[+SEE ALSO?\befibootmgr\b(8), \bupdate-initramfs\b(8), \bbootctl\b(8), \bsystemd-machine-id-setup\b(8).]
'

unset OPT; typeset -A OPT
X="${ print ${USAGE} ; }"
while getopts "${X}" OPT ; do
	case ${OPT} in
		h) showUsage ; exit 0 ;;
		T)	if [[ ${OPTARG} == 'ALL' ]]; then
				typeset -ft ${ typeset +f ; }
			else
				typeset -ft ${OPTARG//,/ }
			fi
			;;
		F) typeset +f && exit 0 ;;
		D) OPT['DEPMOD']=1 ;;
		a) OPT['ADD']+=" ${OPTARG}" ;;
		d) OPT['DEL']+=" ${OPTARG}" ;;
		e) OPT['EFI']="${OPTARG}" ;;
		k) OPT['KERNEL']="${OPTARG}" ;;
		n) OPT['DRY']=1 ;;
		o) OPT['BOPT']+=" ${OPTARG}" ;;
		t) OPT['TRY']=${OPTARG} ;;
		v) OPT['VERB']=1 ;;
		*) showUsage 1 ; exit 1 ;;
	esac
done

X=$((OPTIND-1))
shift $X && OPTIND=1
unset X

doMain "$@"