#!/usr/bin/env bash
set -e
shopt -s extglob

die() { echo "error: $*" >&1; exit 1; }

efi_stub=/usr/lib/systemd/boot/efi/linuxx64.efi.stub
sb_cert=/etc/private/efi/DatabaseKey.crt
sb_key=/etc/private/efi/DatabaseKey.key
in_kernels=()
in_ucode=(/boot/intel-ucode.img)
out_dir=/boot/EFI/Linux

if [[ $1 == @(--help|-h) ]]; then
	echo "Usage: $0 [<vmlinuz_path>...]"
	exit
fi

if (( UID > 0 )); then
	die "Root privileges are needed"
fi

if [[ $1 == --hook ]]; then
	while read -r line; do
		if [[ ${line#/} == @(boot/vmlinuz*|usr/lib/modules/*/vmlinuz) ]]; then
			echo "Signing updated kernel image '$line'."
			in_kernels+=("/${line#/}")
		fi
	done
elif (( $# )); then
	in_kernels=("$@")
fi

if (( ! ${#in_kernels[@]} )); then
	echo "Signing all found kernel images."
	in_kernels=(/boot/vmlinuz-linux!(*.*))
fi

if (( ! ${#in_kernels[@]} )); then
	die "No kernel images specified"
elif [[ ! -e $efi_stub ]]; then
	die "Bootloader stub '$efi_stub' not found"
elif [[ ! -e $sb_key ]]; then
	die "Secure Boot signing key '$sb_key' not found"
elif [[ ! -e $sb_cert ]]; then
	die "Secure Boot signing certificate '$sb_key' not found"
fi

mkdir -p "$out_dir"

for in_vmlinuz in "${in_kernels[@]}"; do
	if [[ ! -s $in_vmlinuz ]]; then
		die "Kernel image '$in_vmlinuz' not found"
	fi

	# Try to guess the pkg_name without invoking a slow `pacman -Qo`
	case $in_vmlinuz in
		/usr/lib/modules/*/vmlinuz)
			pkg_name=$(< "${in_vmlinuz%/*}/pkgbase");;
		/boot/vmlinuz-*)
			pkg_name=${in_vmlinuz##*/vmlinuz-};;
		*)
			die "Could not determine pkg_name name for '$in_vmlinuz'";;
	esac
	pkg_version=$(expac %v "$pkg_name")
	echo "Found kernel image $in_vmlinuz, version $pkg_version ($pkg_name)"

	# .initrd
	in_initrd="/boot/initramfs-${pkg_name}.img"
	if [[ ! -s $in_initrd ]]; then
		die "Initramfs archive '$in_initrd' not found for this kernel"
	fi
	in_initrds=("${in_ucode[@]}" "$in_initrd")

	# .cmdline
	cmdline=$(grep '^[^#]' /etc/kernel/cmdline | tr -s '\n' ' ')

	# .osrel
	os_id=$(. /etc/os-release && echo "$ID")
	os_name=$(. /etc/os-release && echo "$PRETTY_NAME")
	os_release=(
		ID		"$os_id"
		PRETTY_NAME	"$os_name"
		VERSION_ID	"${pkg_version}-signed"
	)

	# Put temporary files where space is plentiful, as osslsigncode
	# silently eats -ENOSPC.
	tmp_dir=$(mktemp -d /tmp/sign-kernels.XXXXXXXX)
	tmp_raw_image=$tmp_dir/bundle_raw.efi
	tmp_signed_image=$tmp_dir/bundle_signed.efi

	echo " - Generating image..."
	objcopy \
		--add-section		.osrel=<(printf '%s="%s"\n' "${os_release[@]}") \
		--change-section-vma	.osrel=0x20000 \
		--add-section		.cmdline=<(printf '%s' "$cmdline") \
		--change-section-vma	.cmdline=0x30000 \
		--add-section		.splash=/usr/share/systemd/bootctl/splash-arch.bmp \
		--change-section-vma	.splash=0x40000 \
		--add-section		.linux="$in_vmlinuz" \
		--change-section-vma	.linux=0x2000000 \
		--add-section		.initrd=<(cat "${in_initrds[@]}") \
		--change-section-vma	.initrd=0x3000000 \
		"$efi_stub" "$tmp_raw_image"

	echo " - Signing bundle (with key $sb_key)..."
	#sbsign \
	#	--cert "$sb_cert" \
	#	--key "$sb_key" \
	#	--output "$tmp_signed_image" \
	#	"$tmp_raw_image"
	osslsigncode sign \
		-certs "$sb_cert" \
		-key "$sb_key" \
		-h sha256 \
		-in "$tmp_raw_image" \
		-out "$tmp_signed_image" \
		> /dev/null

	# Atomically replace the old image with the new one (as far as FAT32
	# allows), and make sure it's flushed to disk in case we're running
	# with root=/bin/sh.
	case $pkg_name in
		linux)	out_image="$out_dir/$os_id.efi";;
		*)	out_image="$out_dir/$os_id-${pkg_name#linux-}.efi";;
	esac
	mv "$tmp_signed_image" "${out_image}.new"
	sync "${out_image}.new"
	mv -b "${out_image}.new" "$out_image"
	sync "$out_image"
	rm -f "$out_image~"
	rm -rf "$tmp_dir"
	echo " - Image $out_image built."
done
