#!/usr/bin/bash

set -eu

usage() {
    echo "Usage:"
    echo "  aboot-deploy [OPTION...] ABOOT_IMG"
    echo
    echo "Options:"
    echo "  -c,--config PATH      - Config files to read env vars from"
    echo "  -d,--destination PATH - Write image to this target (not in --root)"
    echo "  -l,--local            - Flash image to correct slot"
    echo "  -o,--options          - cmdline of next boot, used by OSTree"
    echo "  -r,--root PATH        - Location for the root directory"
    echo
}

get_karg_val() {
    cmdline="$1"
    key="$2"

    echo "$cmdline" | sed "s/\s/\n/g" | grep -m1 "^$key=" | sed "s/^$key=//g"
}

prepare_active_slot() {
    if [ "$NEW_SLOT" = "_a" ]; then
        SLOT_NUM="0"
    elif [ "$NEW_SLOT" = "_b" ]; then
        SLOT_NUM="1"
    else
        echo "No valid new slot in prepare_active_slot: '$NEW_SLOT'"
        exit 12
    fi

    if [ -n "$PREPARE_SWITCH_CMD" ]; then
        $PREPARE_SWITCH_CMD "$SLOT_NUM"
    fi
}

set_active_slot() {
    if [ "$NEW_SLOT" = "_a" ]; then
        LINK="$LINK_A"
        SLOT_NUM="0"
    elif [ "$NEW_SLOT" = "_b" ]; then
        LINK="$LINK_B"
        SLOT_NUM="1"
    else
        echo "No valid new slot in set_active_slot: '$NEW_SLOT'"
        exit 6
    fi

    if ! "$LOCAL"; then
        PREFIX="/ostree/deploy"
        REALPATH=$(realpath "${ROOTDIR}${OSTREE}" | sed "s#.*$PREFIX#deploy#")
        if [ -e "$LINK" ]; then
            rm "$LINK"
        fi

        ln -s "$REALPATH" "$LINK"
        chcon -h system_u:object_r:root_t:s0 "$LINK"
        if $IS_OSBUILD; then
            exit 0 # Appropriate flags set by osbuild
        fi
    fi

    if [ -n "$FINALIZE_SWITCH_CMD" ]; then
        $FINALIZE_SWITCH_CMD "$SLOT_NUM"
    fi
}

FINALIZE_SWITCH_CMD=
PREPARE_SWITCH_CMD=
ABOOT_IMG=
DESTINATION=
VBMETA_DESTINATION=
OPTIONS=
ROOTDIR=
ABOOT_CFG=/boot/aboot.cfg
LOCAL="false"
PARTITION_A=
PARTITION_B=
VBMETA_PARTITION_A=
VBMETA_PARTITION_B=
BOOT_TYPE=

echo "$0 $@"

while [[ $# -gt 0 ]]; do
  case $1 in
    -c|--config)
      ABOOT_CFG="$2"
      shift 2
      ;;
    -d|--destination)
      DESTINATION="$2"
      shift 2
      ;;
    -l|--local)
      LOCAL="true"
      shift 1
      ;;
    -o|--options)
      OPTIONS="$2"
      shift 2
      ;;
    -r|--root)
      ROOTDIR="$2"
      shift 2
      ;;
    -v|--vbmeta-destination)
      VBMETA_DESTINATION="$2"
      shift 2
      ;;
    -*|--*)
      usage
      echo "Unknown option $1"
      exit 1
      ;;
    *)
      break;
      ;;
  esac
done

if [[ $# -lt 1 ]]; then
    usage
    echo "error: No aboot image specified"
    exit 2
fi

ABOOT_IMG="$ROOTDIR/boot/$1"
ABOOT_CFG="$ROOTDIR/$ABOOT_CFG"

if [ -f "$ABOOT_CFG" ]; then
    source "$ABOOT_CFG"
fi

# Compute (possibly per-boot-type) defaults:
BOOTLABEL=boot
if [ "$BOOT_TYPE" == "ukiboot" ]; then
    BOOTLABEL=ukiboot
fi
if [ -z "$PARTITION_A" ]; then
    PARTITION_A="/dev/disk/by-partlabel/${BOOTLABEL}_a"
fi
if [ -z "$PARTITION_B" ]; then
    PARTITION_B="/dev/disk/by-partlabel/${BOOTLABEL}_b"
fi
VBMETA_PARTITION_A="/dev/disk/by-partlabel/vbmeta_a"
VBMETA_PARTITION_B="/dev/disk/by-partlabel/vbmeta_b"

if ! $LOCAL; then
    OSTREE=$(get_karg_val "$OPTIONS" "ostree")
    if ! $LOCAL && [[ $OSTREE != /ostree* ]]; then
        echo "No valid ostree target in '$OPTIONS'"
        exit 3
    fi
fi

CMDLINE="$(cat /proc/cmdline)"
SLOT=$(get_karg_val "$CMDLINE" "androidboot.slot_suffix")
LINK_A="$ROOTDIR/ostree/root.a"
LINK_B="$ROOTDIR/ostree/root.b"
NEW_SLOT=
IS_OSBUILD="false"

if ! $LOCAL && [ ! -e "$LINK_A" ] && [ ! -e "$LINK_B" ]; then
    echo "No symlinks found, assuming osbuild"
    IS_OSBUILD="true"
    NEW_SLOT="_a"
    set_active_slot
    exit 0
fi

SINGLE_SLOT_MODE="false"
if [ "$BOOT_TYPE" == "ukiboot" ]; then
    PREPARE_SWITCH_CMD="ukibootctl prepare-switch"
    FINALIZE_SWITCH_CMD="ukibootctl finalize-switch"
elif command -v abctl; then
    FINALIZE_SWITCH_CMD="abctl --set_active"
elif command -v qbootctl; then
    FINALIZE_SWITCH_CMD="qbootctl -s"
else
    SINGLE_SLOT_MODE="true"
    echo "Warning: No valid ab switching executable, proceeding in single-slot"
    echo "         mode with slot a."
fi

if $SINGLE_SLOT_MODE || [ "$SLOT" = "_b" ]; then
    NEW_SLOT="_a"
elif [ "$SLOT" = "_a" ]; then
    NEW_SLOT="_b"
else
    echo "No valid slot in: '$CMDLINE' and not in single-slot mode"
    exit 4
fi

if [ -z "$DESTINATION" ]; then
    if [ "$NEW_SLOT" = "_a" ]; then
        DESTINATION="$PARTITION_A"
    elif [ "$NEW_SLOT" = "_b" ]; then
        DESTINATION="$PARTITION_B"
    else
        echo "No valid new slot: '$NEW_SLOT'"
        exit 8
    fi
fi

if [ -z "$VBMETA_DESTINATION" ]; then
    if [ "$NEW_SLOT" = "_a" ]; then
        VBMETA_DESTINATION="$VBMETA_PARTITION_A"
    elif [ "$NEW_SLOT" = "_b" ]; then
        VBMETA_DESTINATION="$VBMETA_PARTITION_B"
    else
        echo "No valid new slot: '$NEW_SLOT'"
        exit 7
    fi
fi

if ! [ -e "$DESTINATION" ]; then
    # In osbuild we have to write manually another way, this script is for
    # runtime during actual use on the deployed device
    echo "No destination, not writing via dd, devices mounted:"

    set +e
    ls -ltr /dev/disk/*/* # list mounted devices, to help debug this failure
    set -e

    exit 5
fi

ABOOT_IMG_SIZE=$(stat --format="%s" $ABOOT_IMG)
DEST_SIZE=$(blockdev --getsize64 $DESTINATION)
if [ "$ABOOT_IMG_SIZE" -gt "$DEST_SIZE" ]; then
    echo "Android Boot Image file too large, size: '$ABOOT_IMG_SIZE' dest size: '$DEST_SIZE'"
    exit 9
fi

# Deploy vbmeta only if the image file exists
VBMETA_IMG="${ABOOT_IMG/"aboot-"/"vbmeta-"}"
if ! [ -f "$VBMETA_IMG" ]; then
    VBMETA_IMG=""
fi

if [ -n "$VBMETA_IMG" ]; then
    if ! [ -e "$VBMETA_DESTINATION" ]; then
        # In osbuild we have to write manually another way, this script is for
        # runtime during actual use on the deployed device
        echo "No vbmeta destination, not writing via dd, devices mounted:"

        set +e
        ls -ltr /dev/disk/*/* # list mounted devices, to help debug this failure
        set -e

        exit 10
    fi

    VBMETA_IMG_SIZE=$(stat --format="%s" $VBMETA_IMG)
    VBMETA_DEST_SIZE=$(blockdev --getsize64 $VBMETA_DESTINATION)
    if [ "$VBMETA_IMG_SIZE" -gt "$VBMETA_DEST_SIZE" ]; then
        echo "vbmeta file too large, size: '$VBMETA_IMG_SIZE' dest size: '$VBMETA_DEST_SIZE'"
        exit 11
    fi
fi

prepare_active_slot

# Theoretically we could dd for both partitions in parallel as an optimization,
# and do one sync for all at the end, but this isn't a performance critical
# path and would prefer things to be simple, less error-prone and deterministic.
dd if=$ABOOT_IMG of=$DESTINATION status=progress conv=fsync &> /dev/null
if [ -n "$VBMETA_IMG" ]; then
    dd if=$VBMETA_IMG of=$VBMETA_DESTINATION status=progress conv=fsync &> /dev/null
fi

set_active_slot
