#!/bin/sh
# SPDX-FileCopyrightText: 2023-2024 Helmut Grohne <helmut@subdivi.de>
# SPDX-License-Identifier: MIT
#
# Reimplementation of a subset of systemd's ukify. It is vaguely
# API-compatible. Your mileage will vary. Known differences:
#  * Lots of missing options:
#    * Commands other than build
#    * Secure boot related options
#    * Configuration file support
#  * Additional options --verbose and --deb-arch
#  * --stub accepts a directory containing stubs and selects the right one.

set -eu

die() {
	echo "$*" >&2
	exit 1
}
usage() {
	die "usage: $0 [build] <--linux=|--initrd=|--stub=> [--cmdline=|--deb-arch=|--devicetree=|--efi-arch=|--os-release=|--output=|--section=|--verbose]"
}
usage_error() {
	echo "error: $*" >&2
	usage
}

DEBARCH=
BOOTSTUB=/usr/lib/systemd/boot/efi
OUTPUT=
VERBOSE=

TDIR=
cleanup() {
	if test -n "$TDIR"; then
		rm -Rf "$TDIR"
	fi
}
cleanup_abort() {
	cleanup
	echo aborted >&2
	exit 2
}
trap cleanup EXIT
trap cleanup_abort HUP INT TERM QUIT

# The TDIR will contain one file per section. It can be a symbolic link or
# regular file. The section name is the filename with a leading dot.
TDIR=$(mktemp --directory --tmpdir mini-ukify.XXXXXXXXXX)

add_section_content() {
	test -e "$TDIR/$1" && die "section '$1' already defined"
	printf "%s" "$2" > "$TDIR/$1"
}
add_section_file() {
	test -e "$TDIR/$1" && die "section '$1' already defined"
	test -f "$2" || die "file '$2' does not exist"
	if test "${2#/}" = "$2"; then
		ln --symbolic --relative "$2" "$TDIR/$1"
	else
		ln --symbolic "$2" "$TDIR/$1"
	fi
}
add_section_either() {
	if test "${2#@}" = "$2"; then
		add_section_content "$1" "$2"
	else
		add_section_file "$1" "${2#@}"
	fi
}

opt_cmdline() {
	add_section_either cmdline "$1"
}
opt_deb_arch() {
	case "$1" in
		amd64)
			EFIARCH=x64
			GNU_TYPE=x86_64-linux-gnu
		;;
		arm64)
			EFIARCH=aa64
			GNU_TYPE=aarch64-linux-gnu
		;;
		armhf)
			EFIARCH=arm
			GNU_TYPE=arm-linux-gnueabihf
		;;
		i386)
			EFIARCH=ia32
			GNU_TYPE=i686-linux-gnu
		;;
		riscv64)
			EFIARCH=riscv64
			GNU_TYPE=riscv64-linux-gnu
		;;
		*)
			die "unsupported Debian architecture: $1"
		;;
	esac
	DEBARCH="$1"
}
opt_devicetree() {
	add_section_file dtb "$1"
}
opt_efi_arch() {
	case "$1" in
		aa64) opt_deb_arch arm64 ;;
		arm) opt_deb_arch armhf ;;
		ia32) opt_deb_arch i386 ;;
		x64) opt_deb_arch amd64 ;;
		riscv64) opt_deb_arch "$1" ;;
		*)
			die "unsupported EFI architecture: $1"
		;;
	esac
}
opt_initrd() {
	add_section_file initrd "$1"
}
opt_linux() {
	add_section_file linux "$1"
	test -z "$OUTPUT" && OUTPUT="$1.unsigned.efi"
}
opt_os_release() {
	add_section_either osrel "$1"
}
opt_output() {
	OUTPUT="$1"
}
opt_section() {
	case "$1" in
		*[!a-z]*:*)
			die "invalid section name '${1%%:*}'"
		;;
		*:*)
		;;
		*)
			die "missing section name separated by colon in '$1'"
		;;
	esac
	add_section_either "${1%%:*}" "${1#*:}"
}
opt_stub() {
	BOOTSTUB="$1"
}
opt_verbose() {
	VERBOSE=yes
}
positional=1
positional_1() {
	test "$1" = "build" || usage_error "command not understood: '$1'"
}
positional_2() {
	usage_error "too many positional arguments"
}

while test "$#" -gt 0; do
	case "$1" in
		--cmdline=*|--deb-arch=*|--devicetree=*|--efi-arch=*|--initrd=*|--linux=*|--os-release=*|--output=*|--section=*|--stub=*)
			optname="${1%%=*}"
			optname="${optname#--}"
			test "${optname#*-}" = "$optname" ||
				optname="${optname%%-*}_${optname#*-}"
			"opt_$optname" "${1#*=}"
		;;
		--cmdline|--deb-arch|--devicetree|--efi-arch|--initrd|--linux|--os-release|--output|--section|--stub)
			test "$#" -ge 2 || usage_error "missing argument for $1"
			optname="${1#--}"
			test "${optname#*-}" = "$optname" ||
				optname="${optname%%-*}_${optname#*-}"
			"opt_$optname" "$2"
			shift
		;;
		--verbose)
			"opt_${1#--}"
		;;
		--*)
			usage_error "unrecognized option $1"
		;;
		*)
			"positional_$positional" "$1"
			positional=$((positional + 1))
		;;
	esac
	shift
done

test -z "$DEBARCH" && opt_deb_arch "$(dpkg --print-architecture)"
test -e "$TDIR/linux" || usage_error "missing --linux argument"
test -e "$BOOTSTUB" || die "stub image or directory '$BOOTSTUB' does not exist"
test -e "$OUTPUT" && die "output '$OUTPUT' already exists"

test -d "$BOOTSTUB" && BOOTSTUB="$BOOTSTUB/linux$EFIARCH.efi.stub"
test -f "$BOOTSTUB" || die "efi boot stub $BOOTSTUB not found"

GNU_PREFIX="$GNU_TYPE-"
test "$(dpkg-query --showformat='${db:Status-Status}' --show binutils-multiarch)" = installed &&
	GNU_PREFIX=

# Compute the next multiple of $2 greater than or equal to $1.
align_size() {
	echo "$((($1) + ($2) - 1 - (($1) + ($2) - 1) % ($2)))"
}
alignment=$("${GNU_PREFIX}objdump" --private-headers "$BOOTSTUB" | sed 's/^SectionAlignment\s\+\([0-9]\)/0x/;t;d')
test -z "$alignment" &&
	die "failed to discover the alignment of the efi stub"
test -n "$VERBOSE" &&
	echo "determined efi vma alignment as $alignment"

# Discover the last section in terms of vma + size. We'll append sections
# beyond.
lastoffset=0
# shellcheck disable=SC2034  # unused variables serve documentation
lastoffset="$("${GNU_PREFIX}objdump" --section-headers "$BOOTSTUB" \
  | while read -r idx name size vma lma fileoff algn behind; do
    test -z "$behind" -a "${algn#"2**"}" != "$algn" || continue
    offset=$((0x$vma + 0x$size))
    test "$offset" -gt "$lastoffset" || continue
    lastoffset="$offset"
    echo "$lastoffset"
  done | tail -n1)"
lastoffset=$(align_size "$lastoffset" "$alignment")
test -n "$VERBOSE" &&
	echo "determined minimum efi vma offset as $lastoffset"

# Compute the objdump invocation that constructs the UKI. Successively add
# sections and compute non-overlapping VMAs.
set -- \
	"${GNU_PREFIX}objcopy" \
	--enable-deterministic-archives
for sectionfile in "$TDIR/"*; do
	section="${sectionfile##*/}"
	size=$(stat -Lc%s "$sectionfile")
	size=$(align_size "$size" "$alignment")
	set -- "$@" \
		--add-section ".$section=$sectionfile" \
		--change-section-vma ".$section=$lastoffset"
	lastoffset=$((lastoffset + size))
done
set -- "$@" "$BOOTSTUB" "$OUTPUT"
if test -n "$VERBOSE"; then
	ls -lA "$TDIR"
	printf "%s\n" "$*"
fi
"$@" || die "failed to construct UKI"
