#!/bin/bash

set -eux

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
}

CMDLINE=
INITRD=
KERNEL=
DESTINATION_BOOT=
DESTINATION_VBMETA=
ROOTDIR=
PREPTREE="false"
BLS=

echo "$0 $@"

while [[ $# -gt 0 ]]; do
  case $1 in
    -b|--bls)
      BLS="$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
      ;;
    -*|--*)
      usage
      echo "Unknown option $1"
      exit 1
      ;;
    *)
      break;
      ;;
  esac
done

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

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=

BLS_FILE=(/run/osbuild/tree/boot/loader.1/entries/ostree-1-*.conf)
OSTREE_BOOT_DIR=
if [ -e "$BLS_FILE" ]; then
    ABOOT_CONFIG=$ROOTDIR/$(grep "^abootcfg " $BLS_FILE | sed "s/^abootcfg //g")
    OSTREE_BOOT_DIR="$(dirname $ABOOT_CONFIG)"
fi

if [ -e "$ABOOT_CONFIG" ]; then
    source "$ABOOT_CONFIG"
fi

KERNEL_ADDR=$(( $BASE + $KERNEL_OFFSET ))
RAMDISK_ADDR=$(( $BASE + $RAMDISK_OFFSET ))
TAGS_ADDR=$(( $BASE + $TAGS_OFFSET ))
SECOND_ADDR=$(( $BASE + $SECOND_OFFSET ))
DTB_ADDR=$(( $BASE + $DTB_OFFSET ))

PAGESIZE_ARGS="-c pagesize=$PAGESIZE"
ADDR_ARGS="-c kerneladdr=$KERNEL_ADDR -c ramdiskaddr=$RAMDISK_ADDR -c tagsaddr=$TAGS_ADDR -c secondaddr=$SECOND_ADDR"

if [ -e "$BLS_FILE" ]; then
    CMDLINE="$(grep '^options ' $BLS_FILE | sed 's/^options //g')"
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/$(grep "^fdtdir " $BLS_FILE | sed "s/^fdtdir //g")
    fi

    DTB_PATH="$FDT_DIR/$DTB_FILE"
fi

DTB_ARGS="-d $DTB_PATH -c dtbaddr=$DTB_ADDR"

if [ -e "$BLS_FILE" ]; then
    INITRD=$BOOT_DIR/$(grep "^initrd " $BLS_FILE | sed "s/^initrd //g")
fi

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

if [ -e "$BLS_FILE" ]; then
    KERNEL=$BOOT_DIR/$(grep "^linux " $BLS_FILE | sed "s/^linux //g")
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"
    for FILE in $INITRD*; do
        INITRD="$FILE"
        break
    done

    for FILE in $KERNEL*; do
        KERNEL="$FILE"
        break
    done
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/$(grep "^aboot " $BLS_FILE | sed "s/^aboot //g")
        DESTINATION_VBMETA=$BOOT_DIR/$(grep "^aboot " $BLS_FILE | sed "s/^aboot //g" | sed "s#/aboot-#/vbmeta-#g")
    fi
fi

if command -v mkbootimg; then
    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"
elif command -v abootimg; then
    abootimg --create "$DESTINATION_BOOT" \
      $PAGESIZE_ARGS $ADDR_ARGS \
      -r $INITRD -k $KERNEL \
      -c "cmdline = $CMDLINE" $DTB_ARGS
else
    echo "No valid Android Boot Image creating tool" 
    exit 2
fi

VBMETA="false"
if $VBMETA; then
    # Just create a throwaway certificate for now
    openssl genrsa -out vbmeta.pem 4096
    avbtool extract_public_key --key vbmeta.pem --output vbmeta.avbpubkey
    avbtool make_vbmeta_image --output "$DESTINATION_VBMETA" \
      --padding_size 4096 \
      --include_descriptors_from_image "$DESTINATION_BOOT" \
      --key vbmeta.avbpubkey
fi

if [ "$MAKE_LINK" == "yes" ]; then
    # Update hard link to latest version
    ln -f $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
        # link in /run/osbuild/inputs/tree/boot basically
        ln -f $DESTINATION_VBMETA $ROOTDIR/../tree/boot/ || true
    fi
fi

