#!/bin/sh

set -ue

STATE="$1"
FILES="$2"

if which jq >/dev/null 2>&1; then
    JQ_AVAILABLE=1
else
    JQ_AVAILABLE=0
fi

if which mender-flash >/dev/null 2>&1; then
    MENDER_FLASH_AVAILABLE=1
else
    MENDER_FLASH_AVAILABLE=0
fi

if command -v grub-mender-grubenv-print > /dev/null; then
    PRINTENV=grub-mender-grubenv-print
    SETENV=grub-mender-grubenv-set
else
    PRINTENV=fw_printenv
    SETENV=fw_setenv
fi

resolve_rootfs() {
    case "$1" in
        /dev/root|/dev/disk/by-partlabel/*|/dev/disk/by-partuuid/*)
            # This is a symlink that points to the regular device
            # (e.g. /dev/disk/by-partuuid/b3c4f349-1180-45e1-9a3d-0a6697f4960e --> /dev/sda2)
            readlink -f "$1"
            ;;

        *)
            # Keep any other path as-is
            # (cf. https://github.com/mendersoftware/mender/pull/1613#discussion_r1584353642 for the reasoning)
            echo "$1"
            ;;
    esac
}

parse_conf_file() {
    MENDER_ROOTFS_PART_A=""
    MENDER_ROOTFS_PART_B=""
    # Try first the fallback config file, which has least precedence
    for CONF_FILE in \
            ${MENDER_DATASTORE_DIR:-/var/lib/mender}/mender.conf \
            ${MENDER_CONF_DIR:-/etc/mender}/mender.conf \
    ; do
        if [ "$JQ_AVAILABLE" = 1 ]; then
            # Use the alternative operator "//" to set tmp to "" instead of "null"
            tmp="$(jq -r '.RootfsPartA // empty' < "$CONF_FILE" || true)"
            MENDER_ROOTFS_PART_A="${tmp:-${MENDER_ROOTFS_PART_A}}"
            tmp="$(jq -r '.RootfsPartB // empty' < "$CONF_FILE" || true)"
            MENDER_ROOTFS_PART_B="${tmp:-${MENDER_ROOTFS_PART_B}}"
        else
            # Fall back to line based parsing. Vulnerable to weird JSON nesting, as well as unexpected
            # newlines, although it is unlikely with a regular configuration file.

            # Poor man's case insensitive match.
            MATCH="[Rr][Oo][Oo][Tt][Ff][Ss][Pp][Aa][Rr][Tt][Aa]"
            tmp="$(sed -ne '/"'"$MATCH"'" *: *"[^"]*"/ { s/.*"'"$MATCH"'" *: *"\([^"]*\)".*/\1/; p }' "$CONF_FILE" || true)"
            MENDER_ROOTFS_PART_A="${tmp:-${MENDER_ROOTFS_PART_A}}"
            MATCH="[Rr][Oo][Oo][Tt][Ff][Ss][Pp][Aa][Rr][Tt][Bb]"
            tmp="$(sed -ne '/"'"$MATCH"'" *: *"[^"]*"/ { s/.*"'"$MATCH"'" *: *"\([^"]*\)".*/\1/; p }' "$CONF_FILE" || true)"
            MENDER_ROOTFS_PART_B="${tmp:-${MENDER_ROOTFS_PART_B}}"
        fi
    done

    if [ -z "$MENDER_ROOTFS_PART_A" ] || [ -z "$MENDER_ROOTFS_PART_B" ]; then
        echo "Cannot parse RootfsPartA/B in any configuration file!" 1>&2
        return 1
    fi

    # For UBI, standardize on the `/dev/` variant. The kernel only accepts an argument without
    # `/dev/`, but all userspace tools use the `/dev/` variant.
    MENDER_ROOTFS_PART_A="$(echo "$MENDER_ROOTFS_PART_A" | sed -e 's,^ubi,/dev/ubi,')"
    MENDER_ROOTFS_PART_B="$(echo "$MENDER_ROOTFS_PART_B" | sed -e 's,^ubi,/dev/ubi,')"

    # Resolve paths if required.
    MENDER_ROOTFS_PART_A="$(resolve_rootfs "$MENDER_ROOTFS_PART_A")"
    MENDER_ROOTFS_PART_B="$(resolve_rootfs "$MENDER_ROOTFS_PART_B")"

    # Extract the partition number from the regular device path (e.g. /dev/sda2 --> 2).
    MENDER_ROOTFS_PART_A_NUMBER="$(echo "$MENDER_ROOTFS_PART_A" | grep -Eo '[0-9]+$' || true)"
    MENDER_ROOTFS_PART_B_NUMBER="$(echo "$MENDER_ROOTFS_PART_B" | grep -Eo '[0-9]+$' || true)"

    return 0
}

set_upgrade_vars() {
    active_num="$(${PRINTENV} mender_boot_part)"
    active_num="${active_num#mender_boot_part=}"
    if test "$active_num" -eq "$MENDER_ROOTFS_PART_A_NUMBER"; then
        active=$MENDER_ROOTFS_PART_A
        passive=$MENDER_ROOTFS_PART_B
        passive_num=$MENDER_ROOTFS_PART_B_NUMBER
    else
        active=$MENDER_ROOTFS_PART_B
        passive=$MENDER_ROOTFS_PART_A
        passive_num=$MENDER_ROOTFS_PART_A_NUMBER
    fi
    active_num_hex=$(printf '%x' "$active_num")
    passive_num_hex=$(printf '%x' "$passive_num")
    upgrade_available="$(${PRINTENV} upgrade_available)"
    upgrade_available="${upgrade_available#upgrade_available=}"
}

check_environment_canary() {
    mender_check_saveenv_canary="$(${PRINTENV} mender_check_saveenv_canary)"
    if [ "$mender_check_saveenv_canary" = "mender_check_saveenv_canary=1" ]; then
        # If the check canary exists (added during build), we need to check the real canary to make
        # sure that the boot loader was successful in adding it during boot.
        mender_saveenv_canary="$(${PRINTENV} mender_saveenv_canary)"
        if [ "$mender_saveenv_canary" != "mender_saveenv_canary=1" ]; then
            cat 1>&2 <<'EOF'
`mender_check_saveenv_canary` was set in the boot environment, but
`mender_saveenv_canary` was not. This is an indication that the bootloader
integration is not working correctly, and the bootloader was not able to save
an environment which we can read from userspace. Please re-check your
bootloader integration, and refer to the section on Bootloader support in the
Mender documentation if you need more information.
EOF
            return 1
        fi
    fi
    return 0
}

check_device_matches_root() {
    case "$1" in
        /dev/ubi*)
            # UBI is a bit peculiar. Trying to take the device number of the root device does not
            # match the major/minor number of the UBI device file, even though that filesystem is
            # mounted. So fall back on name comparison for UBI.

            # Standardize on the `/dev/` variant. The kernel only accepts an argument without
            # `/dev/`, but all userspace tools use the `/dev/` variant.
            ROOT_DEVICE="$(mount | grep -F ' on / ' | sed -e 's/ .*//; s,^ubi,/dev/ubi,')"
            if [ "$1" = "$ROOT_DEVICE" ]; then
                return 0
            else
                echo "Mounted root ($ROOT_DEVICE) does not match boot loader environment ($1)!" 1>&2
                return 1
            fi
            ;;
        *)
            # Match major/minor device number against mounted root device.
            if [ "$(stat -L -c %02t%02T "$1")" = "$(stat -L -c %04D /)" ]; then
                return 0
            else
                echo "Mounted root does not match boot loader environment ($1)!" 1>&2
                return 1
            fi
            ;;
    esac
}

check_requirements() {
    parse_conf_file
    check_environment_canary
    set_upgrade_vars
}

case "$STATE" in
    ProvidePayloadFileSizes)
        echo "Yes"
        ;;

    Download)
        echo "This module supports DownloadWithFileSizes only" 1>&2
        exit 1
        ;;

    DownloadWithFileSizes)
        check_requirements

        if [ "$upgrade_available" != 0 ]; then
            echo "Unexpected \`upgrade_available=$upgrade_available\` in $STATE." 1>&2
            exit 1
        fi
        check_device_matches_root "$active"

        line="$(cat stream-next)"
        file="$(echo "$line" | cut -d' ' -f1)"
        size="$(echo "$line" | cut -d' ' -f2)"
        if [ -z "$file" ] || [ -z "$size" ]; then
            echo "Cannot parse line from stream-next, got: $line" 1>&2
            exit 1
        fi
        if [ "$MENDER_FLASH_AVAILABLE" = 1 ]; then
            mender-flash --input-size "$size" --input "$file" --output "$passive"
        elif echo "$passive" | grep "^/dev/ubi" > /dev/null; then
            ubiupdatevol "$passive" --size="$size" "$file"
        else
            cat "$file" > "$passive"
            sync
        fi
        if [ "$(cat stream-next)" != "" ]; then
            echo "More than one file in payload" 1>&2
            exit 1
        fi
        ;;

    ArtifactInstall)
        check_requirements

        cat > "$FILES/tmp/orig-part.tmp" <<EOF
orig_part_num=$active_num
orig_part_num_hex=$active_num_hex
EOF
        sync "$FILES/tmp/orig-part.tmp"
        mv "$FILES/tmp/orig-part.tmp" "$FILES/tmp/orig-part"
        sync "$FILES/tmp"

        if [ "$upgrade_available" != 0 ]; then
            echo "Unexpected \`upgrade_available=$upgrade_available\` in $STATE." 1>&2
            exit 1
        fi
        check_device_matches_root "$active"

        ${SETENV} -s - <<EOF
mender_boot_part=$passive_num
mender_boot_part_hex=$passive_num_hex
upgrade_available=1
bootcount=0
EOF
        ;;

    NeedsArtifactReboot)
        echo "Automatic"
        ;;

    SupportsRollback)
        echo "Yes"
        ;;

    ArtifactVerifyReboot)
        check_requirements
        if test "$upgrade_available" != 1; then
            exit 1
        fi
        check_device_matches_root "$active"
        ;;

    ArtifactVerifyRollbackReboot)
        check_requirements
        if test "$upgrade_available" = 1; then
            exit 1
        fi
        check_device_matches_root "$active"
        ;;

    ArtifactCommit)
        check_requirements

        if [ "$upgrade_available" != 1 ]; then
            echo "Unexpected \`upgrade_available=$upgrade_available\` in $STATE." 1>&2

            # If we get here, an upgrade in standalone mode failed to boot and the user is trying to commit from the old OS.
            # This communicates to the user that the upgrade failed.
            echo "Upgrade failed and was reverted: refusing to commit!" 1>&2
            exit 1
        fi
        check_device_matches_root "$active"

        ${SETENV} upgrade_available 0
        ;;

    ArtifactRollback)
        # If we cannot parse the config file, exit anyway and let the bootloader handle the rollback
        parse_conf_file || exit 0
        check_requirements

        # We do not use `check_device_matches_root` here, since we can be on either partition at
        # this point.

        if test "$upgrade_available" = 1; then
            # If upgrade_available = 1, then we know that the bootloader will roll back for us, even
            # if we fail here.
            ${SETENV} -s - <<EOF || true
mender_boot_part=$passive_num
mender_boot_part_hex=$passive_num_hex
upgrade_available=0
EOF
        elif [ -f "$FILES/tmp/orig-part" ]; then
            . "$FILES/tmp/orig-part"
            ${SETENV} -s - <<EOF
mender_boot_part=$orig_part_num
mender_boot_part_hex=$orig_part_num_hex
upgrade_available=0
EOF
        fi
        ;;
esac
exit 0
