#!/usr/bin/env bash
# Themes the GRUB passphrase prompt of an encrypted install.
#
# When /boot lives inside the LUKS volume, GRUB itself asks for the passphrase
# before it can read grub.cfg, so the boot menu theme cannot apply: the prompt
# is drawn as plain text in the top left corner. The prompt comes from the core
# image, so the only way to style it is to rebuild that image with a font, the
# artwork and the translations embedded.
#
# grub-install recreates the core image on every grub upgrade and overwrites
# this one, which is why the hook calls this script again afterwards.
#
# Nothing here is worth an unbootable machine. The image is built and verified
# in a temporary directory, and the one on the ESP is only replaced once that
# copy proves usable. Any doubt and the script gives up quietly, leaving the
# stock GRUB prompt in place: plain, but working.
set -uo pipefail

# A way to keep this out of the way while building an image or debugging a
# machine, where rewriting the host bootloader is never what is wanted.
[[ -e /etc/grub-biglinux-hooks.disabled ]] && exit 0

# Marks our own images so the script can tell them from the stock ones.
signature=biglinux-crypt-prompt

# Giving up is the safe outcome, so it is not a failure for the transaction.
bail() {
	printf '%s\n' "grub-theme-biglinux-crypt: $1, keeping the stock GRUB prompt" >&2
	exit 0
}

theme_directory=/boot/grub/themes/biglinux
artwork=$theme_directory/crypt-background.png
font=/usr/share/grub/ter-u20n.pf2
esp=/boot/efi

for tool in grub-mkimage grub-script-check cryptsetup msgunfmt msgfmt gettext findmnt tar awk; do
	command -v "$tool" >/dev/null || bail "$tool is unavailable"
done

# A BIOS boot has no EFI image to replace.
[[ -d /sys/firmware/efi ]] || exit 0

# The prompt only exists when GRUB has to unlock the root device itself, and
# the UUID has to come from that device: other LUKS volumes may be attached.
root_source=$(findmnt -no SOURCE /)
# On Btrfs findmnt appends the subvolume, as in /dev/mapper/luks-x[/@].
root_source=${root_source%%[*}
root_mapper=${root_source##*/}
[[ -e /dev/mapper/$root_mapper ]] || exit 0
backing_device=$(cryptsetup status "$root_mapper" 2>/dev/null |
	awk '$1 == "device:" { print $2; exit }')
[[ -n $backing_device ]] || exit 0
crypt_uuid=$(cryptsetup luksUUID "$backing_device" 2>/dev/null | tr -d -)
[[ -n $crypt_uuid ]] || bail "no LUKS UUID found for $backing_device"

# The image points GRUB at /boot/grub inside the encrypted root, so that has to
# be where /boot/grub really is. It is not when /boot is a filesystem or a
# subvolume of its own, and it is not inside an installer chroot, where findmnt
# answers about the host instead of the system being installed. Both would
# unlock the disk and then drop to the grub> prompt, which is the failure this
# whole script exists to avoid.
[[ $(findmnt -no SOURCE --target "$theme_directory" 2>/dev/null) == "$(findmnt -no SOURCE --target / 2>/dev/null)" ]] ||
	bail "$theme_directory is not on the root filesystem the image would point at"

# A signed bootloader must not be replaced by an image we cannot sign, so this
# has to fail closed. Asking mokutil did the opposite: it is not installed by
# default, and its absence read as "Secure Boot is off". The firmware itself
# always answers - four bytes of EFI attributes, then the flag.
secure_boot=unknown
for variable in /sys/firmware/efi/efivars/SecureBoot-*; do
	[[ -r $variable ]] || continue
	case $(od -An -tu1 -j4 -N1 "$variable" | tr -d '[:space:]') in
	0) secure_boot=disabled ;;
	1) secure_boot=enabled ;;
	esac
	break
done
# sbctl resigns in its own hook, which runs after this one.
if [[ $secure_boot != disabled ]] && ! command -v sbctl >/dev/null; then
	bail "Secure Boot is $secure_boot and no signing tool is available"
fi

[[ -f $artwork ]] || bail "$artwork is missing"
[[ -f $font ]] || bail "$font is missing"

# Pick the directory holding our bootloader: an ESP shared with Windows also
# has Microsoft/, and the order find returns is not alphabetical.
candidates=()
while IFS= read -r candidate; do
	[[ -f $candidate ]] || continue
	grep -aqF grub "$candidate" || continue
	candidates+=("$candidate")
done < <(find "$esp/EFI" -mindepth 2 -maxdepth 2 -name 'grubx64.efi' 2>/dev/null | sort)

case ${#candidates[@]} in
0) bail 'no GRUB EFI binary was found on the ESP' ;;
1) target=${candidates[0]} ;;
*)
	# An ESP shared with another Linux holds that system's GRUB too, and it
	# sorts ahead of ours as often as not. Replacing it would point its boot
	# entry at our encrypted root. Ours is the one that already knows this
	# volume: grub-install embeds the LUKS UUID it has to unlock, and after
	# the first run our own signature is there as well.
	target=
	for candidate in "${candidates[@]}"; do
		if grep -aqF "$signature" "$candidate" || grep -aqF "$crypt_uuid" "$candidate"; then
			target=$candidate
			break
		fi
	done
	[[ -n $target ]] ||
		bail 'the ESP holds several GRUB binaries and none of them is identifiably ours'
	;;
esac

# On Btrfs the subvolume is part of the path GRUB has to reach after unlocking.
subvolume=$(findmnt -no OPTIONS --target / | grep -o 'subvol=[^,]*' | cut -d= -f2)
prefix="(cryptouuid/$crypt_uuid)${subvolume:-}/boot/grub"

export TEXTDOMAINDIR=/usr/share/locale
export TEXTDOMAIN=grub-theme-biglinux
# The prompt belongs to the machine, not to whoever happens to run pacman: an
# upgrade over SSH carries that person's LANG, which may be another language.
language=$(awk -F= '$1 == "LANG" { gsub(/"/, "", $2); print $2; exit }' /etc/locale.conf 2>/dev/null)
language=${language:-${LANG:-}}
language=${language%%.*}
[[ -n $language ]] || language=en
export LANG=$language.UTF-8

# The boot menu has GRUB's unicode font loaded next to the theme's and falls
# back glyph by glyph. This prompt has only what is packed into it, so a
# language Terminus cannot draw would face a row of empty boxes - worse than
# reading the message in English. Terminus does cover Latin, Cyrillic, Greek
# and Hebrew; it has no CJK.
case $language in
ja | ko | zh*)
	if [[ -f /usr/share/grub/unicode.pf2 ]]; then
		font=/usr/share/grub/unicode.pf2
	else
		# Better plain English in a font that draws than boxes.
		export LANG=en.UTF-8
		language=en
	fi
	;;
esac

workspace=$(mktemp -d) || bail 'no temporary directory available'
trap 'rm -rf "$workspace"' EXIT
mkdir -p "$workspace/locale" || bail 'the temporary directory is not writable'
cp "$font" "$workspace/font.pf2" || bail 'the font could not be copied'
cp "$artwork" "$workspace/background.png" || bail 'the artwork could not be copied'

# GRUB ships catalogs for its own strings. Reuse the one matching the system
# and rewrite the passphrase prompt: its default spells out the disk UUID,
# which means nothing to the person typing the password.
catalog=/boot/grub/locale/$language.mo
# GRUB has no pt_BR for every language it has a pt for, and none at all for
# English, where its own strings are already the source text.
[[ -f $catalog ]] || catalog=/boot/grub/locale/${language%%_*}.mo
if [[ -f $catalog ]] && msgunfmt "$catalog" >"$workspace/catalog.po"; then
	# Replace the translation in place. The c-format flag has to go with it:
	# it would make msgfmt reject a translation that drops the arguments.
	awk -v replacement="$(gettext -- 'Encryption password: ')" '
		/^#, c-format$/ { flag = $0; next }
		/^msgid "Enter passphrase for/ {
			print; print "msgstr \"" replacement "\""; flag = ""; inside = 1; next
		}
		inside && NF == 0 { inside = 0; print; next }
		inside { next }
		{ if (flag != "") { print flag; flag = "" } print }
	' "$workspace/catalog.po" >"$workspace/prompt.po"
else
	# Without a catalog GRUB would print its own prompt, which spells out the
	# device and its UUID and overruns the card. One entry is enough to
	# replace it; the rest of GRUB stays in its original language.
	{
		printf '%s\n' 'msgid ""' 'msgstr ""' \
			'"Content-Type: text/plain; charset=UTF-8\\n"' ''
		printf '%s\n' 'msgid "Enter passphrase for %s%s%s (%s): "'
		printf 'msgstr "%s"\n' "$(gettext -- 'Encryption password: ')"
	} >"$workspace/prompt.po"
fi
msgfmt -o "$workspace/locale/$language.mo" "$workspace/prompt.po" ||
	bail 'the translated prompt could not be compiled'

# Terminus is 10 pixels wide at this size and the prompt is drawn 1024 pixels
# wide, so centering is a matter of counting characters.
center() {
	local text=$1 length
	length=$(printf '%s' "$text" | LC_ALL=C.UTF-8 wc -m)
	printf '%*s%s' $(((102 - length) / 2)) '' "$text"
}

# GRUB prints the prompt itself, right after this indentation, so the text
# belongs to the catalog above and must not be repeated here.
prompt_indent=$(center "$(gettext -- 'Encryption password: ')")
prompt_indent=${prompt_indent%%[! ]*}

# $root and $prefix belong to GRUB and are expanded at boot, not by the shell.
# shellcheck disable=SC2016
{
	printf '# %s\n' "$signature"
	printf '%s\n' 'set root=(memdisk)' 'set prefix=($root)/'
	printf 'insmod %s\n' all_video efi_gop efi_uga video_bochs video_cirrus video \
		gfxterm gfxterm_background font png gettext
	printf '%s\n' 'loadfont /font.pf2' 'set locale_dir=($root)/locale' \
		"set lang=$language" 'set gfxmode=1024x768,auto' 'terminal_output gfxterm' \
		'background_image -m stretch /background.png'
	printf '%s\n' 'echo ""' 'echo ""' 'echo ""' 'echo ""' 'echo ""' 'echo ""' 'echo ""' \
		'echo ""' 'echo ""' 'echo ""' 'echo ""' 'echo ""' 'echo ""'
	printf 'echo "%s"\n' "$(center "$(gettext -- 'This disk is protected by encryption')")"
	printf '%s\n' 'echo ""'
	printf 'echo "%s"\n' "$(center "$(gettext -- 'Type the disk password to continue')")"
	printf '%s\n' 'echo ""'
	printf 'echo "%s"\n' "$(center "$(gettext -- 'It does not appear while you type')")"
	printf '%s\n' 'echo ""'
	printf 'echo -n "%s"\n' "$prompt_indent"
	# Unlock this volume only: -a would also prompt for unrelated LUKS disks.
	# Clearing the artwork lets the boot menu render with its own theme.
	# The prefix needs the device in parentheses, exactly like grub.cfg writes
	# it: root holds a bare name, so "$root/..." would not resolve and GRUB
	# would unlock the disk and then drop to its command prompt.
	printf '%s\n' "cryptomount -u $crypt_uuid" 'background_image' 'clear' \
		"set root=cryptouuid/$crypt_uuid" "set prefix=(\$root)${subvolume:-}/boot/grub" \
		'insmod normal' 'normal'
} >"$workspace/early-grub.cfg" || bail 'the early configuration could not be written'

# GRUB parses this before any module is available; a syntax error there is an
# unbootable machine.
grub-script-check "$workspace/early-grub.cfg" >/dev/null 2>&1 ||
	bail 'the generated early configuration is not valid GRUB script'

tar cf "$workspace/memdisk.tar" -C "$workspace" font.pf2 background.png locale ||
	bail 'the memdisk could not be packed'

# Everything used before /boot is reachable has to live in the image: the
# graphics stack for the prompt, the crypto stack for LUKS2 with Argon2id and
# every filesystem the installer can put /boot on.
modules=(
	memdisk tar configfile normal search search_fs_uuid part_gpt part_msdos
	btrfs ext2 xfs f2fs fat cryptodisk luks luks2 argon2 pbkdf2
	gcry_rijndael gcry_sha256 gcry_sha512
	all_video efi_gop efi_uga video_bochs video_cirrus video
	gfxterm gfxterm_background font png gettext
	echo test loadenv linux halt reboot minicmd
)

# The list above covers the defaults the installer produces. A volume built
# with another cipher or hash needs its own module, and without it GRUB rejects
# the right passphrase forever - the stock image avoids this by probing the
# device, so this has to as well.
dump=$(cryptsetup luksDump "$backing_device" 2>/dev/null)
for algorithm in serpent:gcry_serpent twofish:gcry_twofish camellia:gcry_camellia \
	sha1:gcry_sha1 ripemd160:gcry_rmd160 whirlpool:gcry_whirlpool; do
	[[ $dump == *"${algorithm%%:*}"* ]] && modules+=("${algorithm##*:}")
done

# GRUB reads the early configuration before it can load anything, so every
# module it names has to be inside the image already. An insmod of a module
# that was left out stops the boot exactly where the passphrase is asked.
while read -r module; do
	[[ " ${modules[*]} " == *" $module "* ]] ||
		bail "the early configuration loads $module, which is not embedded in the image"
done < <(awk '$1 == "insmod" { print $2 }' "$workspace/early-grub.cfg")

grub-mkimage -c "$workspace/early-grub.cfg" -o "$workspace/grubx64.efi" \
	-O x86_64-efi -m "$workspace/memdisk.tar" -p "$prefix" \
	"${modules[@]}" 2>"$workspace/mkimage.log" ||
	bail "grub-mkimage failed: $(tr '\n' ' ' <"$workspace/mkimage.log")"

# An image that cannot unlock the disk is worse than no theme: check that the
# crypto path and the prefix really made it in before trusting this file.
#
# The prefix is the one that matters most. It has to name the device in
# parentheses, exactly as grub.cfg writes it, because GRUB drops the
# parentheses when assigning to a variable - the earlier version of this script
# wrote a bare $root, and the machine unlocked the disk and then sat at the
# grub> prompt. Reading it back out of the built image is the only check that
# proves what GRUB will actually see.
for marker in cryptodisk luks2 argon2 btrfs ext2 "$crypt_uuid" "$signature" "$prefix"; do
	grep -aqF "$marker" "$workspace/grubx64.efi" ||
		bail "the generated image is missing $marker"
done
image_size=$(stat -c '%s' "$workspace/grubx64.efi")
((image_size > 400000)) || bail 'the generated image is implausibly small'

# Everything above happened in a temporary directory. What follows touches the
# firmware's own copy of the bootloader, so it is the only part that can leave
# the machine unable to start.
#
# The new image is written beside the target and put in place by rename, so the
# file the firmware loads is either the old one or the new one and never half of
# either. The ESP is FAT and often the 100 MB one Windows created, so the order
# matters: writing the backup first is what could fill the filesystem and then
# truncate the very loader the backup was meant to protect.
install_image() {
	local destination=$1 staged=$1.biglinux-new
	cp -f "$workspace/grubx64.efi" "$staged" || return 1
	# Keep the stock image as the way back, but never overwrite that backup
	# with one of ours: after the first run the target is already themed.
	if ! grep -aqF "$signature" "$destination"; then
		cp -f "$destination" "$destination.pre-biglinux" || {
			rm -f "$staged"
			return 1
		}
	fi
	# FAT has no journal, so the bytes have to be on the disk before the name
	# points at them.
	sync "$staged"
	mv -f "$staged" "$destination" || {
		rm -f "$staged"
		return 1
	}
	sync "${destination%/*}"
}

install_image "$target" || {
	printf '%s\n' "grub-theme-biglinux-crypt: the themed image could not be installed, $target is unchanged" >&2
	exit 1
}

# The removable fallback is only ours to replace when it is also GRUB. Its
# failure is not fatal: the machine still boots through the entry above.
fallback=$esp/EFI/boot/bootx64.efi
if [[ -f $fallback ]] && grep -aqF grub "$fallback"; then
	install_image "$fallback" ||
		printf '%s\n' "grub-theme-biglinux-crypt: $fallback is unchanged" >&2
fi

printf '%s\n' "grub-theme-biglinux-crypt: encryption prompt written to $target"
