#!/bin/sh
#
# Purpose {{{
# This script will create an Unified Kernel + UEFI entry for ArchLinux
#   1. Verify dependancies (efibootmgr, objcopy…).
#   2. Clean older unified kernels.
#   3. Backup last unified kernel.
#   4. Create a new unified kernel from vmlinuz, initramfs, UUID of root,…
#   5. Ensure an UEFI entry exists for this unified kernel.
#
# 2023-08-17
# }}}
# TODO {{{
# 1. Encrypted LVM is not yet supported. See crypt_part_uuid relative comments.
# }}}
# Flags {{{
## Exit on error {{{
set -o errexit
## }}}
## Exit on unset var {{{
### Use "${VARNAME-}" to test a var that may not have been set
set -o nounset
## }}}
## Pipeline command is treated as failed {{{
### Not available in POSIX sh − https://github.com/koalaman/shellcheck/wiki/SC3040
#set -o pipefail
## }}}
## Help with debugging {{{
### Call the script by prefixing it with "TRACE=1 ./script.sh"
if [ "${TRACE-0}" -eq 1 ]; then set -o xtrace; fi
## }}}
# }}}
# Vars {{{
PROGNAME=$(basename "${0}"); readonly PROGNAME
PROGDIR=$(readlink --canonicalize-missing $(dirname "${0}")); readonly PROGDIR
ARGS="${*}"; readonly ARGS
readonly NBARGS="${#}"
[ -z "${DEBUG-}" ] && DEBUG=1
## Export DEBUG for sub-script
export DEBUG

## Default values for Unified kernel
readonly ROOT_UUID_DEFAULT=$(findmnt --kernel --noheadings --output UUID -- /)
readonly ROOT_FSTYPE_DEFAULT=$(findmnt --kernel --noheadings --output FSTYPE -- /)
readonly CRYPT_PART_UUID_DEFAULT=$(blkid | sed --silent 's;/dev/.*: UUID="\(.*\)" TYPE="crypto_LUKS".*;\1;p')
readonly UEFI_BOOT_STUB_FILE_DEFAULT="/usr/lib/systemd/boot/efi/linuxx64.efi.stub"

## Default values for UEFI entry
readonly EFI_BASE_LABEL_DEFAULT="Arch Linux unified"
readonly EFI_MOUNT_PATH_DEFAULT="/boot/efi"

## Temp files
readonly temp_kernel_command_file="/tmp/kernel.command.temp"

## Colors
readonly PURPLE='\033[1;35m'
readonly RED='\033[0;31m'
readonly RESET='\033[0m'
readonly COLOR_DEBUG="${PURPLE}"
# }}}
usage() {                                                       # {{{

	cat <<- HELP
usage: $PROGNAME [-d|-h|-b|-f|-k|-u|-l|-m]

Script to build unified kernel for Arch Linux will all required informations
(vmlinuz, initramfs, UUID of root's partition,…) and create an UEFI entry.

EXAMPLES :
    - Build unified kernel with default :
        ${PROGNAME}

    - Build unified kernel with current command line :
        ${PROGNAME} --kernel "\$(< /proc/cmdline)"

OPTIONS :
    -d,--debug
        Enable debug messages.

    -h,--help
        Print this help message.

Unified kernel OPTIONS :

    -b,--boot-stub,--uefi-boot
        Define UEFI boot stub file to use.
        (default: ${UEFI_BOOT_STUB_FILE_DEFAULT})
        For more infos, see: https://man.archlinux.org/man/linuxx64.efi.stub.7.en

    -f,--fs,--root-fs
        Define a different root partition's filesystem.
        (default: ${ROOT_FSTYPE_DEFAULT})
        Thanks to command : findmnt --kernel --noheadings --output FSTYPE -- /

    -u,--uuid,--root-uuid
        Define a different root partition's UUID.
        (default: ${ROOT_UUID_DEFAULT})
        Thanks to command : findmnt --kernel --noheadings --output UUID -- /

    -c,--crypt,--crypt-uuid
        Define a different LUKS partition's UUID.
	(default: ${CRYPT_PART_UUID_DEFAULT})
        Thanks to command : blkid | sed --silent 's;/dev/.*crypto_LUKS.*UUID="\(.*\)";\1;p'

    -k,--kernel,--kernel-cmd
        Define the kernel command line to use.
        (default: build with root-fs|root-uuid + ro options)

UEFI entry OPTIONS :

    -l,--label
        Set a new label for UEFI entry.
        (default: ${EFI_BASE_LABEL_DEFAULT})

    -m,--mount-path,--esp
        Define a different mountpoint for the EFI partition.
        (default: ${EFI_MOUNT_PATH_DEFAULT})
HELP
}
# }}}
debug_message() {                                               # {{{

	local_debug_message="${1}"

	## Print message if DEBUG is enable (=0)
	[ "${DEBUG}" -eq "0" ] && printf '\e[1;35m%-6b\e[m\n' "DEBUG − ${PROGNAME} : ${local_debug_message}"

	unset local_debug_message

	return 0
}
# }}}
error_message() {                                               # {{{

	local_error_message="${1}"
	local_error_code="${2}"

	## Print message
	printf '%b\n' "ERROR − ${PROGNAME} : ${RED}${local_error_message}${RESET}" >&2

	unset local_error_message

	exit "${local_error_code:=66}"
}
# }}}
define_vars() {                                                 # {{{

	# If root_uuid wasn't defined (argument) {{{
	if [ -z "${root_uuid-}" ]; then
		## Use default value
		readonly root_uuid="${ROOT_UUID_DEFAULT}"
	fi
	# }}}
	# If root_fstype wasn't defined (argument) {{{
	if [ -z "${root_fstype-}" ]; then
		## Use default value
		readonly root_fstype="${ROOT_FSTYPE_DEFAULT}"
	fi
	# }}}
	# If crypt_part_uuid wasn't defined (argument) {{{
	if [ -z "${crypt_part_uuid-}" ]; then
		## Use default value
		readonly crypt_part_uuid="${CRYPT_PART_UUID_DEFAULT}"
	fi
	# }}}
	# If a LUKS partition exists AND if kernel_command_line wasn't defined (argument) {{{
	if $(blkid | grep --silent -- "LUKS") && [ -z "${kernel_command_line-}" ]; then
		## Use default value for encrypted LVM
		readonly kernel_command_line="root=UUID=${root_uuid} rootfstype=${root_fstype} ro cryptdevice=UUID=${crypt_part_uuid}:lvm"
	fi
	# }}}
	# If kernel_command_line wasn't defined (argument) {{{
	if [ -z "${kernel_command_line-}" ]; then
		## Use default value
		readonly kernel_command_line="root=UUID=${root_uuid} rootfstype=${root_fstype} ro"
		## For encrypted LVM
		#readonly kernel_command_line="root=UUID=${root_uuid} rootfstype=${root_fstype} ro cryptdevice=UUID=${crypt_part_uuid}:lvm"
	fi
	# }}}
	# If uefi_boot_stub_file wasn't defined (argument) {{{
	if [ -z "${uefi_boot_stub_file-}" ]; then
		## Use default value
		readonly uefi_boot_stub_file="${UEFI_BOOT_STUB_FILE_DEFAULT}"
	fi
	# }}}

	# If efi_base_label wasn't defined (argument) {{{
	if [ -z "${efi_base_label-}" ]; then
		## Use default value
		readonly efi_base_label="${EFI_BASE_LABEL_DEFAULT}"
	fi
	# }}}
	# If efi_mount_path wasn't defined (argument) {{{
	if [ -z "${efi_mount_path-}" ]; then
		## Use default value
		readonly efi_mount_path="${EFI_MOUNT_PATH_DEFAULT}"
	fi
	# }}}

	# If esp_device wasn't defined (argument) {{{
	if [ -z "${esp_device-}" ]; then
		## Use default value
		readonly esp_device=$(findmnt -kno SOURCE "${efi_mount_path}" | grep --invert-match -- systemd | sed s-/dev/--)
	fi
	# }}}
	# If esp_disk wasn't defined (argument) {{{
	if [ -z "${esp_disk-}" ]; then
		## Use default value
		readonly esp_disk=$(lsblk /dev/"${esp_device}" -no pkname)
	fi
	# }}}
	# If esp_part wasn't defined (argument) {{{
	if [ -z "${esp_part-}" ]; then
		## Use default value
		readonly esp_part=$(cat /sys/class/block/"${esp_device}"/partition)
	fi
	# }}}

}
# }}}

is_command_available() {                                        # {{{

	local_command_available_cmd="${1}"
	debug_prefix="${2:-}"

	## Return False by default
	return_command_available="1"

	if [ "$(command -v ${local_command_available_cmd})" ]; then
		debug_message "${debug_prefix}is_command_available − \
${RED}${local_command_available_cmd}${COLOR_DEBUG} seems present on this host."
		return_command_available="0"
	else
		debug_message "${debug_prefix}is_command_available − \
${RED}${local_command_available_cmd}${COLOR_DEBUG} is not available on this host."
		return_command_available="1"
	fi

	unset local_command_available_cmd
	unset debug_prefix

	return "${return_command_available}"
}
# }}}

main() {                                                        # {{{

	debug_message "--- MAIN BEGIN"

	## If a efibootmgr is missing {{{
	### Exit with error message
	is_command_available "efibootmgr" "| " \
		|| error_message "No efibootmgr command available. Please install efibootmgr package with your package manager (pacman -S efibootmgr)." 01
	## }}}
	## If a objcopy is missing {{{
	### Exit with error message
	is_command_available "objcopy" "| " \
		|| error_message "No objcopy command available. Please install binutils package with your package manager (pacman -S binutils)." 02
	## }}}

	## Define all vars
	define_vars
	debug_message "| Define vars"

	## If UEFI boot stub is missing {{{
	### Exit with error message
	test -f "${uefi_boot_stub_file}" \
		|| error_message "No UEFI boot stub file available (${uefi_boot_stub_file}). Please verify your systemd installation (systemd or systemd-boot-efi packages)." 03
	## }}}
	## Ensure EFI device is mounted {{{
	if ! mountpoint --quiet "${efi_mount_path}"; then
		mount "${efi_mount_path}" \
			|| error_message "Can't mount EFI device" 04
	fi
	## }}}

	## Put Kernel command line in temp file {{{
	rm --force -- "${temp_kernel_command_file}" ; touch "${temp_kernel_command_file}"
	printf "%s" "${kernel_command_line}" >> "${temp_kernel_command_file}"
	## }}}

	## Calculate address values to use for each section {{{
	osrel_offs=$(objdump --section-headers -- "${uefi_boot_stub_file}" | awk 'NF==7 {size=strtonum("0x"$3); offset=strtonum("0x"$4)} END {print size + offset}')
	cmdline_offs=$((osrel_offs + $(stat --dereference --format=%s -- "/usr/lib/os-release")))
	linux_offs=$((cmdline_offs + $(stat --dereference --format=%s -- "${temp_kernel_command_file}")))
	#linux_offs=$((cmdline_offs + $(stat --dereference --format=%s -- "/proc/cmdline")))
	initrd_offs=$((linux_offs + $(stat --dereference --format=%s -- "/boot/vmlinuz-linux")))
	## }}}

	## Debug message {{{
	debug_message "Try to build Unified kernel with from this informations :\r
  * UEFI boot stub file: ${RED}${uefi_boot_stub_file}${COLOR_DEBUG}\r
  * Root's UUID: ${RED}${root_uuid}${COLOR_DEBUG}\r
  * Root's filesystem: ${RED}${root_fstype}${COLOR_DEBUG}\r
  * Kernel command line: ${RED}${kernel_command_line}${COLOR_DEBUG}\r
  * EFI device: ${RED}${esp_device}${COLOR_DEBUG}\r
  * EFI disk: ${RED}${esp_disk}${COLOR_DEBUG}\r
  * EFI partition number: ${RED}${esp_part}${COLOR_DEBUG}\r
Addresses values :\r
  * osrel : ${RED}${osrel_offs}${COLOR_DEBUG}\r
  * cmdline : ${RED}${cmdline_offs}${COLOR_DEBUG}\r
  * linux : ${RED}${linux_offs}${COLOR_DEBUG}\r
  * initrd : ${RED}${initrd_offs}${COLOR_DEBUG}\r
And generate UEFI entry :\r
  * Label: ${RED}${efi_base_label}${COLOR_DEBUG}"

	## }}}

	## Ensure EFI mount path subdirectories exists
	mkdir --parents -- "${efi_mount_path}/EFI/arch/"
	## Create unified kernel blob {{{
	objcopy \
		--add-section .osrel="/usr/lib/os-release" --change-section-vma .osrel=$(printf 0x%x $osrel_offs) \
		--add-section .cmdline="${temp_kernel_command_file}" --change-section-vma .cmdline=$(printf 0x%x $cmdline_offs) \
		--add-section .linux="/boot/vmlinuz-linux" --change-section-vma .linux=$(printf 0x%x $linux_offs) \
		--add-section .initrd="/boot/initramfs-linux.img" --change-section-vma .initrd=$(printf 0x%x $initrd_offs) \
		"${uefi_boot_stub_file}" "${efi_mount_path}/EFI/arch/linux.arch.efi"
	## }}}
	## If no UEFI entry for this label + kernel {{{
	## Create one
	efibootmgr | grep --quiet -- "${efi_base_label}.*\\\EFI\\\arch\\\linux.arch.efi" \
		|| efibootmgr --disk /dev/"${esp_disk}" --part "${esp_part}" --create --label "${efi_base_label}" --loader "\\EFI\\arch\\linux.arch.efi"
	## }}}

	## Remove temp files
	rm --force -- "${temp_kernel_command_file}"

	debug_message "--- MAIN END"
}
# }}}

# Manage arguments                                                # {{{
# This code can't be in a function due to argument management

if [ ! "${NBARGS}" -eq "0" ]; then

	manage_arg="0"

	## If the first argument ask for help (h|help|-h|-help|-*h|-*help) {{{
	if printf -- '%s' "${1-}" | grep --quiet --extended-regexp -- "^-*h(elp)?$"; then
		usage
		exit 0
	fi
	## }}}

	## If the first argument is not an option
	if ! printf -- '%s' "${1}" | grep --quiet --extended-regexp -- "^-+";
	then
		## Print help message and exit
		printf '%b\n' "${RED}Invalid option: ${1}${RESET}"
		printf '%b\n' "---"
		usage

		exit 1
	fi

	# Parse all options (start with a "-") one by one
	while printf -- '%s' "${1-}" | grep --quiet --extended-regexp -- "^-+"; do

	case "${1}" in
		## OPTIONS
		-d|--debug )                          ## debug
			DEBUG=0
			debug_message "--- Manage argument BEGIN"
			;;
		-h|--help )                           ## display help message
			usage
			exit 0
			;;
		## Unified kernel OPTIONS
		-u|--boot-stub|--uefi-boot )          ## Define uefi_boot_stub_file with given arg
			## Move to the next argument
			shift
			## Define var
			readonly uefi_boot_stub_file="${1}"
			;;
		-k|--kernel|--kernel-cmd )            ## Define kernel_command_line with given arg
			## Move to the next argument
			shift
			## Define var
			readonly kernel_command_line="${1}"
			;;
		-f|--fs|--root-fs )                   ## Define root_fstype with given arg
			## Move to the next argument
			shift
			## Define var
			readonly root_fstype="${1}"
			;;
		-u|--uuid|--root-uuid )               ## Define root_uuid with given arg
			## Move to the next argument
			shift
			## Define var
			readonly root_uuid="${1}"
			;;
		## UEFI entry OPTIONS
		-l|--label )                          ## Define efi_base_label with given arg
			## Move to the next argument
			shift
			## Define var
			readonly efi_base_label="${1}"
			;;
		-m|--mount-path|--esp )               ## Define efi_mount_path with given arg
			## Move to the next argument
			shift
			## Define var
			readonly efi_mount_path="${1}"
			;;
		* )                                   ## unknow option
			printf '%b\n' "${RED}Invalid option: ${1}${RESET}"
			printf '%b\n' "---"
			usage
			exit 1
			;;
	esac

	debug_message "| ${RED}${1}${COLOR_DEBUG} option managed."

	## Move to the next argument
	shift
	manage_arg=$((manage_arg+1))

	done

	debug_message "| ${RED}${manage_arg}${COLOR_DEBUG} argument(s) successfully managed."
else
	debug_message "| No arguments/options to manage."
fi

	debug_message "--- Manage argument END"
# }}}

main

exit 0