#!/bin/sh
# vim: set ts=4:
#---help---
# Usage: efi-mkuki [options] <vmlinuz> [<microcode>...] [<initrd>]
#        efi-mkuki <-h | -V>
#
# Create an EFI Unified Kernel Image (UKI) - a single EFI PE executable
# combining an EFI stub loader, a kernel image, an initramfs image, the kernel
# command line, and optionally a CPU microcode update image.
#
# Arguments:
#   <vmlinuz>            Location of Linux kernel image file.
#   <microcode>...       Location of microcode file(s) (optional).
#   <initrd>             Location of initramdisk file.
#
# Options:
#   -c <cmdline | file>  Kernel cmdline, or location of file with kernel cmdline
#                        (if begins with "/" or "."). Defaults to /proc/cmdline.
#
#   -k <version | file>  Kernel release ("uname -r" information), or location of
#                        file with kernel release (if begins with "/" or ".").
#                        This is optional on x86 and x86_64 (if not specified,
#                        it's parsed from <vmlinuz>), but required on others.
#
#   -o <file>            Write output into <file>. Defaults to <vmlinuz>.efi.
#
#   -r <file>            Location of osrel file. Defaults to /etc/os-release.
#
#   -s <file>            Location of splash image (optional).
#
#   -S <file>            Location of EFI stub file. Defaults to
#                        linux<march>.efi.stub in /usr/lib/gummiboot where
#                        <march> is UEFI machine type (x64, ia32, aa64, arm).
#
#   -h                   Show this message and exit.
#
#   -V                   Print version.
#
# Please report bugs at <https://github.com/jirutka/efi-mkuki/issues>.
#---help---
set -eu

if ( set -o pipefail 2>/dev/null ); then
	set -o pipefail
fi

PROGNAME='efi-mkuki'
VERSION='1.0.0'
EFISTUB_DIR='/usr/lib/systemd/boot/efi'
ALIGNMENT=4096  # bytes

help() {
	sed -n '/^#---help---/,/^#---help---/p' "$0" | sed 's/^# \?//; 1d;$d;'
}

die() {
	echo "$PROGNAME: $@" >&2
	exit 2
}

# Prints value of the specified variable $1.
getval() {
	eval "printf '%s\n' \"\$$1\""
}

# Calculate the next offset by aligning the size $2 to the $ALIGNMENT boundary
# and adding the current offset $1.
next_offset() {
	printf '%d' "$(( $1 + ( ($2 + $ALIGNMENT - 1) / $ALIGNMENT * $ALIGNMENT) ))"
}

kver_x86() {
	local vmlinuz="$1"
	# https://gitlab.archlinux.org/archlinux/mkinitcpio/mkinitcpio/-/blob/a6ef10d3f4f24029fc649686ba36f692d8575a7a/functions#L142-162
	# https://www.kernel.org/doc/html/v6.7/arch/x86/boot.html
	# https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/arch/x86/boot/header.S?h=v6.7
	local offset; offset="$(od -An -j0x20E -dN2 "$vmlinuz")" || return
	local kver="$(dd if="$vmlinuz" bs=1 count=127 skip="$(($offset + 0x200))" 2>/dev/null)"

	# Read the first word from this string as the kernel version.
	printf '%s' "${kver%% *}"
}


# Defaults
cmdline=/proc/cmdline
output=
osrel=/etc/os-release
splash=/dev/null
efistub=
kver=

while getopts ':c:k:o:r:s:S:hV' OPT; do
	case "$OPT" in
		c) cmdline=$OPTARG;;
		k) kver=$OPTARG;;
		o) output=$OPTARG;;
		r) osrel=$OPTARG;;
		s) splash=$OPTARG;;
		S) efistub=$OPTARG;;
		h) help; exit 0;;
		V) echo "$PROGNAME $VERSION"; exit 0;;
		\?) die "unknown option: -$OPTARG";;
	esac
done
shift $((OPTIND - 1))

[ $# -ge 1 ] || die "invalid number of arguments, see '$PROGNAME -h'"

if ! [ "$efistub" ]; then
	case "$(uname -m)" in
		aarch64) march=aa64;;
		arm*) march=arm;;
		x86 | i686) march=ia32;;
		x86_64) march=x64;;
		*) die "unknown architecture: $(uname -m)";;
	esac
	efistub="$EFISTUB_DIR/linux$march.efi.stub"
fi
[ -f "$efistub" ] || die "EFI stub '$efistub' does not exist!"

tmpdir=$(mktemp -dt $PROGNAME.XXXXXX)
trap "rm -f $tmpdir/*; rmdir $tmpdir" EXIT HUP INT TERM

case "$cmdline" in
	/* | .*) grep '^[^#]' "$cmdline" | tr -s '\n' ' ' > "$tmpdir"/cmdline;;
	*) printf '%s\n' "$cmdline" > "$tmpdir"/cmdline;;
esac
cmdline="$tmpdir/cmdline"

linux=$1; shift

uname="$tmpdir/uname"
case "$kver" in
	'') case "$(uname -m)" in
	    	i686 | x86 | x86_64) kver_x86 "$linux";;
	    	*) die 'missing required option -k';;
	    esac ;;
	/* | .*) head -n 1 "$kver";;
	*) printf '%s\n' "$kver";;
esac > "$uname"

grep -q '^[0-9]\+\.[0-9][0-9A-Za-z_.-]*$' "$uname" \
	|| die "invalid kernel version (-k): $(cat "$uname")"

initrd=${1:-"/dev/null"}
if [ $# -gt 1 ]; then
	initrd="$tmpdir/initrd"
	cat "$@" > "$initrd"
fi

[ "$output" ] || output="$linux.efi"

# Take the size and offset of the last EFI stub section to compute the next
# section offset.
offset="$(objdump -h -w "$efistub" | awk 'END {
	vma = ("0x"$4) + 0;
	size = ("0x"$3) + 0;
	print vma + size
}')"
test "$offset" -eq "$offset" 2>/dev/null \
	|| die 'failed to resolve EFI stub offset!'
offset="$(next_offset 0 "$offset")"  # align offset

objcopy_opts=''
for section in uname osrel cmdline splash linux initrd; do
	filename="$(getval "$section")"
	[ "$filename" = '/dev/null' ] && continue

	objcopy_opts="$objcopy_opts --add-section \".$section=$filename\" --change-section-vma \".$section=$offset\""
	size="$(stat -Lc '%s' "$filename")"
	offset="$(next_offset "$offset" "$size")"
done

eval set -- "$objcopy_opts"
objcopy "$@" "$efistub" "$output"

# This corrects putting sections below the base image of the stub.
objcopy --adjust-vma 0 "$output"
