#!/usr/bin/bash

set -eu

usage() {
	echo "Usage:"
	echo "  aboot-update [OPTION...] KERNEL_VERSION"
	echo
	echo "Options:"
	echo "  -b,--bls FILE                - bls file to process cmdline, initrd, kernel, dtb etc."
	echo "  -c,--cmdline CMDLINE         - kernel cmdline"
	echo "  -i,--initrd INITRD           - initrd/initramfs"
	echo "  -k,--kernel KERNEL           - kernel image"
	echo "  -r,--root PATH               - The location for the root directory where images are stored"
	echo "  -d,--destination PATH        - Write boot image to this file/device (not in --root)"
	echo "  -p,--preptree                - For preptree stage of ostree"
	echo "  -v,--vbmeta-destination PATH - Write vbmeta image to this file/device (not in --root)"
	echo "  -C,--compress-kernel         - Compress the kernel (default: true)"
	echo
}

expand_file_glob()
{
	GLOB=$1
	GLOB_PATH=$(dirname "${GLOB}")
	GLOB_FILE=$(basename "${GLOB}")
	find "${GLOB_PATH}" -maxdepth 1 -name "${GLOB_FILE}-*"
}

extract_cmdline_value() {
	local key=$1
	echo "$CMDLINE" | grep -oP "\s*$key=\K[^ ]+" || true
}

remove_cmdline_key_value() {
	local key=$1
	# shellcheck disable=SC2001
	echo "$CMDLINE" | sed "s/\s*$key=[^ ]* *//g"
}

CMDLINE=
INITRD=
KERNEL=
DESTINATION_BOOT=
DESTINATION_VBMETA=
ROOTDIR=
PREPTREE="false"
COMPRESS_KERNEL="true"
BLS_FILE=
BOOT_TYPE=

echo "$0 $*"

while [[ $# -gt 0 ]]; do
	case $1 in
		-b|--bls)
			BLS_FILE="$2"
			shift 2
			;;
		-c|--cmdline)
			CMDLINE="$2"
			shift 2
			;;
		-i|--initrd)
			INITRD="$2"
			shift 2
			;;
		-k|--kernel)
			KERNEL="$2"
			shift 2
			;;
		-d|--destination)
			DESTINATION_BOOT="$2"
			shift 2
			;;
		-r|--root)
			ROOTDIR="$2"
			shift 2
			;;
		-p|--preptree)
			PREPTREE="true"
			shift 1
			;;
		-v|--vbmeta-destination)
			DESTINATION_VBMETA="$2"
			shift 2
			;;
		-C|--compress-kernel)
			COMPRESS_KERNEL="$2"
			shift 2
			;;
		-*)
			usage
			echo "Unknown option $1"
			exit 1
			;;
		*)
			break;
			;;
	esac
done

if [[ $# -lt 1 ]]; then
	usage
	echo "error: No kernel version specified"
	exit 1
fi

set -x

KERNEL_VERSION="$1"

if $PREPTREE; then
	# The kernel initially is here when we first prepare the OS image, reading
	# the comments in ostree-sysroot-deploy.c might be helpful to explain
	BOOT_DIR="$ROOTDIR/usr/lib/ostree-boot"
else
	BOOT_DIR="$ROOTDIR/boot"
fi

ABOOT_CONFIG="$BOOT_DIR/aboot.cfg"

MAKE_LINK=no
if [ "$DESTINATION_BOOT" == "" ]; then
	DESTINATION_BOOT="$BOOT_DIR/aboot-$KERNEL_VERSION.img"
	DESTINATION_VBMETA="$BOOT_DIR/vbmeta-$KERNEL_VERSION.img"
	MAKE_LINK=yes
fi

# Defaults (from mkbootimg)
PAGESIZE=2048
BASE=0x10000000
KERNEL_OFFSET=0x00008000
RAMDISK_OFFSET=0x01000000
TAGS_OFFSET=0x00000100
SECOND_OFFSET=0x00f00000
DTB_OFFSET=0x01f00000
DTB_FILE=
MKBOOTIMG="true"

if [ "$BLS_FILE" = "" ] ; then
	BLS_FILE=$(find /run/osbuild/tree/boot/loader.1/entries/ostree-*.conf | head -n 1)
	if [ -e "$BLS_FILE" ]; then
		# We should skip mkbootimg in this case, we just need to link to
		# tree/boot
		MKBOOTIMG="false"
	fi
fi

OSTREE_BOOT_DIR=
if [ -e "$BLS_FILE" ]; then
	ABOOT_CONFIG=$ROOTDIR/$(sed -n 's/^abootcfg //pg' "$BLS_FILE")
	OSTREE_BOOT_DIR=$(dirname "$ABOOT_CONFIG")
fi

if [ -e "$ABOOT_CONFIG" ]; then
	# shellcheck source=/dev/null
	source "$ABOOT_CONFIG"
fi

if [ -e "$BLS_FILE" ]; then
	CMDLINE=$(sed -n 's/^options //pg' "$BLS_FILE")
fi

if [ "$CMDLINE" == "" ]; then
	if [[ -f $ROOTDIR/usr/etc/kernel/cmdline ]]; then # ostree
		CMDLINE=$(cat "$ROOTDIR/usr/etc/kernel/cmdline")
	elif [[ -f $ROOTDIR/etc/kernel/cmdline ]]; then
		CMDLINE=$(cat "$ROOTDIR/etc/kernel/cmdline")
	elif [[ -f $ROOTDIR/usr/lib/kernel/cmdline ]]; then
		CMDLINE=$(cat "$ROOTDIR/usr/lib/kernel/cmdline")
	elif [[ -f $ROOTDIR/proc/cmdline ]]; then
		CMDLINE=$(cat "$ROOTDIR/proc/cmdline")
	fi
fi

# If absolute path, just use this, this is useful for dtbs packaged separately
# so they don't have to keep 'uname -r' version in sync, as the location of
# these files change based on that.
if [ "${DTB_FILE:0:1}" = "/" ]; then
	if $PREPTREE || [ -z "$OSTREE_BOOT_DIR" ]; then
		DTB_PATH="$ROOTDIR/$DTB_FILE"
	else
		# This else should only be used in ostree non-preptree case
		DTB_PATH="$OSTREE_BOOT_DIR/../../../$DTB_FILE"
	fi
else
	FDT_DIR=$BOOT_DIR/dtb-$KERNEL_VERSION
	if [ -e "$BLS_FILE" ]; then
		FDT_DIR=$BOOT_DIR/$(sed -n 's/^fdtdir //pg' "$BLS_FILE" )
	fi

	DTB_PATH="$FDT_DIR/$DTB_FILE"
fi

if [ -e "$BLS_FILE" ]; then
	INITRD=$BOOT_DIR/$(sed -n 's/^initrd //pg' "$BLS_FILE")
fi

if [ "$INITRD" == "" ]; then
	INITRD="$BOOT_DIR/initramfs-$KERNEL_VERSION.img"
fi

if [ -e "$BLS_FILE" ]; then
	KERNEL=$BOOT_DIR/$(sed -n 's/^linux //pg' "$BLS_FILE")
fi

if [ "$KERNEL" == "" ]; then
	KERNEL="$BOOT_DIR/vmlinuz-$KERNEL_VERSION"
fi

if $PREPTREE; then
	# this section is sort of similar to get_kernel_from_tree_legacy_layouts
	# in ostree in that we just get the first two files starting with vmlinuz-
	# and initramfs- prefixes, this is needed for the first deploy
	# first deploy files are in locations like:
	# /usr/lib/ostree-boot/vmlinuz-5.14.0-230.194.el9iv.aarch64-217d248ed4e397b0129b14185eff75dc56568d0ef732b6e3a690144eab2a9003
	# /usr/lib/ostree-boot/initramfs-5.14.0-230.194.el9iv.aarch64.img-217d248ed4e397b0129b14185eff75dc56568d0ef732b6e3a690144eab2a9003

	ABOOT_DIR="$ROOTDIR/usr/lib/modules/$KERNEL_VERSION"
	DESTINATION_BOOT="$ABOOT_DIR/aboot-$KERNEL_VERSION.img" # first time called no BLS file
	DESTINATION_VBMETA="$ABOOT_DIR/vbmeta-$KERNEL_VERSION.img"
	INITRD=$(expand_file_glob "${INITRD}")
	KERNEL=$(expand_file_glob "${KERNEL}")
else
	# after the first preptree, you can just use bls entries
	# linux /ostree/centos-8c84cbe2779b25e40bffb8854aa44b3b349c992dcfed62fccd798735cd5ecdf7/vmlinuz-5.14.0-230.194.el9iv.aarch64
	# initrd /ostree/centos-8c84cbe2779b25e40bffb8854aa44b3b349c992dcfed62fccd798735cd5ecdf7/initramfs-5.14.0-230.194.el9iv.aarch64.img

	ABOOT_DIR="$ROOTDIR"
	if [ -e "$BLS_FILE" ]; then
		DESTINATION_BOOT=$BOOT_DIR/$(sed -n 's/^aboot //pg' "$BLS_FILE")
		DESTINATION_VBMETA=$BOOT_DIR/$(sed -n 's/^aboot //pg' "$BLS_FILE" | sed "s#/aboot-#/vbmeta-#g")
	fi
fi

if [[ "$(file "$KERNEL")" =~ "EFI application" ]] ; then
	# The Fedora aarch64 kernel uses the vmlinuz.efi target. ABL won't boot
	# these binaries, so use unzboot (https://github.com/eballetbo/unzboot)
	# to extract the kernel.
	TMP=$(mktemp)
	trap 'rm -f "$TMP"' EXIT
	unzboot "$KERNEL" "$TMP"
	KERNEL="$TMP"

	# The kernel returned by unzboot is not compressed. Compression is
	# enabled here.
	if [[ "$COMPRESS_KERNEL" == "true" ]] ; then
		echo "Compressing kernel..."
		TMP=$(mktemp)
		gzip --no-name --stdout "$KERNEL" > "$TMP"
		mv "$TMP" "$KERNEL"
	fi
else
	# The kernel is gzip compressed. Booting an uncompressed kernel saves
	# ~100ms off the overall boot in the firmware.
	if [[ "$COMPRESS_KERNEL" == "false" ]] ; then
		UNCOMPRESSED_KERNEL=$(mktemp)
		gunzip --stdout "$KERNEL" > "$UNCOMPRESSED_KERNEL"
		echo "Adding uncompressed kernel to the boot image"
		KERNEL="$UNCOMPRESSED_KERNEL"
		trap 'rm -f "$UNCOMPRESSED_KERNEL"' EXIT
	fi

fi

VBMETA="true"

if [ "$BOOT_TYPE" == "ukiboot" ]; then
    MKBOOTIMG=false
    VBMETA=false
    DTB_ARG=""
    if [ "${DTB_FILE}" != "" ]; then
        DTB_ARG="--devicetree=${DTB_PATH}"
    fi
    ukify build \
        --linux="${KERNEL}" \
        --initrd="${INITRD}" \
        --cmdline="${CMDLINE}" \
        ${DTB_ARG} \
        --output="${DESTINATION_BOOT}"
fi

if $MKBOOTIMG; then
	# The Android Boot Image v2 is limited to 511 bytes for the kernel command line.
	# Some kernel parameters are longer, so add them to the bootconfig file if they
	# are present. Note that not all kernel parameters can be added to the
	# bootconfig file, such as the ones that are early parameters, which there are
	# hundreds. So only add certain parameters to the bootconfig file.
	BOOT_CONFIG_ENTRIES=""

	# shellcheck disable=SC2043
	for KEY in systemd.random-seed; do
		VALUE=$(extract_cmdline_value "$KEY")
		if [ "${VALUE}" != "" ] ; then
			BOOT_CONFIG_ENTRIES="$BOOT_CONFIG_ENTRIES$KEY = $VALUE"$'\n'
			CMDLINE=$(remove_cmdline_key_value "$KEY")
		fi
	done

	if [ "$BOOT_CONFIG_ENTRIES" != "" ] ; then
		BOOTCONFIG=$(mktemp)
		echo "kernel {" > "$BOOTCONFIG"
		echo "$BOOT_CONFIG_ENTRIES" >> "$BOOTCONFIG"
		echo "}" >> "$BOOTCONFIG"

		CMDLINE="$CMDLINE bootconfig"
		bootconfig -a "$BOOTCONFIG" "$INITRD"
		rm -f "$BOOTCONFIG"
	fi

	mkbootimg \
		--base "$BASE" \
		--kernel_offset "$KERNEL_OFFSET" \
		--ramdisk_offset "$RAMDISK_OFFSET" \
		--tags_offset "$TAGS_OFFSET" \
		--pagesize "$PAGESIZE" \
		--second_offset "$SECOND_OFFSET" \
		--ramdisk "$INITRD" \
		--dtb_offset "$DTB_OFFSET" \
		--dtb "$DTB_PATH" \
		--kernel "$KERNEL" \
		--cmdline "$CMDLINE" \
		--header_version 2 \
		-o "$DESTINATION_BOOT"

	if $VBMETA; then
		# Just create a throwaway certificate for now
		AVB_PRIVATE_KEY_PEM_FILE="vbmeta.pem"
		openssl genrsa -out $AVB_PRIVATE_KEY_PEM_FILE 4096
		avbtool add_hash_footer --image "$DESTINATION_BOOT" \
			--partition_name "boot" --algorithm SHA256_RSA4096 \
			--key "$AVB_PRIVATE_KEY_PEM_FILE" --rollback_index 0 \
			--dynamic_partition_size
		avbtool make_vbmeta_image --include_descriptors_from_image \
			"$DESTINATION_BOOT" --algorithm SHA256_RSA4096 \
			--key "$AVB_PRIVATE_KEY_PEM_FILE" --rollback_index 0 --output \
			"$DESTINATION_VBMETA"
	fi
fi

if [ "$MAKE_LINK" == "yes" ]; then
	# Update hard link to latest version
	cp "$DESTINATION_BOOT" "$ABOOT_DIR/aboot.img"
	# link in /run/osbuild/inputs/tree/boot basically
	ln -f "$DESTINATION_BOOT" "$ROOTDIR/../tree/boot/" || true

	if $VBMETA; then
		ln -f "$DESTINATION_VBMETA" "$ABOOT_DIR/vbmeta.img" || true
		# link in /run/osbuild/inputs/tree/boot basically
		ln -f "$DESTINATION_VBMETA" "$ROOTDIR/../tree/boot/" || true
	fi
fi

