#!/bin/bash -

# Copyright (C) MOXA Inc. All rights reserved.
# This software is distributed under the terms of the MOXA SOFTWARE NOTICE.
# See the file MOXA-SOFTWARE-NOTICE for details.
#
# Name:
#	MOXA Bootloader Upgrade Utility
#
# Description:
#	Upgrade moxa device's bootloader
#
# Copyright (C) Moxa, Inc. All rights reserved.
# Copyright (C) 2019	Wes Huang	<Wes.Huang@moxa.com>
# Copyright (C) 2021	Remus Wu	<remusty.wu@moxa.com>
# Copyright (C) 2022	Henry LC Chen	<HenryLC.Chen@moxa.com>

BASENAME="mx-bootloader-upgrade-tool"
INTERFACE_NAME="mx-bootloader-mgmt upgrade"

IOTHINX_UBOOT_DEV=/dev/mmcblk2boot0
IOTHINX_FLAG=0
V1200_FLAG=0

STANDARD_UBOOT_DIR="/lib/firmware/moxa/bootloader"

if [ -f "/proc/device-tree/model" ]; then
    device_tree_model_name="$(tr -d '\0' </proc/device-tree/model)"
fi

if echo "${device_tree_model_name}" | grep -q "ioThinx"; then
    DEFAULT_UBOOT_DEV="$IOTHINX_UBOOT_DEV"
    DEFAULT_UBOOT_SIZE=0x100000
    IOTHINX_FLAG=1
else
    DEFAULT_UBOOT_PARTITION_LABEL="u-boot 1"
    DEFAULT_UBOOT_DEV=/dev/$(grep "$DEFAULT_UBOOT_PARTITION_LABEL" </proc/mtd | awk '{print $1}' | sed 's/[[:punct:]]//g')
    DEFAULT_UBOOT_SIZE=0x$(grep "$DEFAULT_UBOOT_PARTITION_LABEL" </proc/mtd | awk '{print $2}' | sed -r 's/0*([0-9])/\1/')
fi

if echo "${device_tree_model_name}" | grep -q "V1200"; then
    V1200_FLAG=1
fi

INFO_LENGTH=1024
UBOOT_LENGTH=$(($((DEFAULT_UBOOT_SIZE)) - INFO_LENGTH))
INFO_START_ADDRESS=0x$(printf "%x\n" $UBOOT_LENGTH)

IOTHINX_INFO_START_ADDRESS=0x1$(printf "%x\n" $UBOOT_LENGTH)

_caculate_sha256sum() {
    local source="$1"
    local length="$2"

    if [ ! -e "$source" ]; then
        echo "$source it not exist"
        exit 1
    fi

    if [ "$source" = "$IOTHINX_UBOOT_DEV" ]; then
        local block_size=1024
        local skip_address=1048576
        dd if="$source" bs=$((1 * "$block_size")) count=$(("$length" / "$block_size")) skip=$(("$skip_address" / "$block_size")) status=none | sha256sum | awk '{print $1}'
    else
        dd if="$source" bs="$length" count=1 status=none | sha256sum | awk '{print $1}'
    fi
}

_caculate_md5sum() {
    local source="$1"
    local length="$2"

    if [ ! -e "$source" ]; then
        echo "$source it not exist"
        exit 1
    fi

    if [ "$source" = "$IOTHINX_UBOOT_DEV" ]; then
        local block_size=1024
        local skip_address=1048576
        dd if="$source" bs=$((1 * "$block_size")) count=$(("$length" / "$block_size")) skip=$(("$skip_address" / "$block_size")) status=none | md5sum | awk '{print $1}'
    else
        dd if="$source" bs="$length" count=1 status=none | md5sum | awk '{print $1}'
    fi
}

_question() {
    local yes="$1"
    local message="$2"
    local choice

    if [ "$yes" != "y" ]; then
        read -r -p "$message" choice
        if [ "${choice,,}" != "y" ]; then
            exit 1
        fi
    else
        echo "${message}${yes}"
    fi
}

_query_info() {
    local source="$1"
    local item="$2"
    local result

    if [ ! -e "$source" ]; then
        $LOGGER_ECHO "$BASENAME" "$source it not exist"
        exit 1
    fi

    if [ -z "$item" ]; then
        $LOGGER_ECHO "$BASENAME" "Please assign query item"
        exit 1
    fi

    # 1024 means 1024 bytes show in 1 line.
    if [ -z "$bootloader_info" ]; then
        if [ "$source" = "$IOTHINX_UBOOT_DEV" ]; then
            bootloader_info=$(hexdump -v -e "$INFO_LENGTH \"%_p\" \"\n\"" "$source" -s "$IOTHINX_INFO_START_ADDRESS" -n "$INFO_LENGTH")
        else
            bootloader_info=$(hexdump -v -e "$INFO_LENGTH \"%_p\" \"\n\"" "$source" -s "$INFO_START_ADDRESS" -n "$INFO_LENGTH")
        fi
    fi

    result=$(grep -o -E "${item}=\S{1,}" <<<"$bootloader_info" | cut -d'=' -f 2)
    if [ -z "$result" ]; then
        $LOGGER_ECHO "$BASENAME" "Can't query $item information"
        exit 1
    fi

    echo "$result"
}

verify_checksum() {
    local source=$1

    if [ "$(_caculate_sha256sum "$source" $UBOOT_LENGTH)" != "$(_query_info "$source" sha256sum)" ]; then
        $LOGGER_ECHO "$BASENAME" "The $source sha256sum is incorrect!"
        return 1
    elif [ "$(_caculate_md5sum "$source" $UBOOT_LENGTH)" != "$(_query_info "$source" md5sum)" ]; then
        $LOGGER_ECHO "$BASENAME" "The $source md5sum is incorrect!"
        return 1
    fi

    return 0
}

verify_modelname_V1200() {
    local EEPROM_I2C_BUS=0
    local EEPROM_I2C_SLAVE_ADDRESS=0x50
    #nHwMajorVersion1   140 (0x8C)
    local EEPROM_HW_MAJOR_VERSION1_OFFSET=0x8c

    # Values of nHwMajorVersion1
    local DDR_MICRON_1=0x00
    local DDR_MICRON_2=0xff
    local DDR_NANYA=0x01

    if ! eeprom_hw_major_version1=$(i2cget -f -y "$EEPROM_I2C_BUS" "$EEPROM_I2C_SLAVE_ADDRESS" "$EEPROM_HW_MAJOR_VERSION1_OFFSET"); then
        $LOGGER_ECHO "$BASENAME" "Failed to read nHwMajorVersion1 from EEPROM, abort."
        exit 1
    fi

    case "${eeprom_hw_major_version1}" in
    "$DDR_MICRON_1" | "$DDR_MICRON_2")
        local TARGET_MODEL="V1200"
        ;;
    "$DDR_NANYA")
        local TARGET_MODEL="V1222"
        ;;
    *)
        $LOGGER_ECHO "$BASENAME" "Invalid nHwMajorVersion1 in EEPROM"
        return 1
        ;;
    esac

    if [ ! "$(_query_info "$source_binary" modelname)" = "$TARGET_MODEL" ]; then
        return 1
    fi
    return 0
}

verify_modelname() {
    local source_binary=$1

    if [ "$V1200_FLAG" = "1" ]; then
        if ! verify_modelname_V1200; then
            $LOGGER_ECHO "$BASENAME" "The bootloader is invalid for current model, abort."
            return 1
        fi
    fi

    return 0
}

upgrade() {
    local source_binary=$1
    local target_device=$DEFAULT_UBOOT_DEV
    local source_binary_biosver
    local target_device_biosver
    local systemfailback

    if ! verify_checksum "$source_binary"; then
        exit 1
    fi

    if ! verify_modelname "$source_binary"; then
        exit 1
    fi

    source_binary_biosver=$(_query_info "$source_binary" biosver)
    target_device_biosver=$(_query_info "$target_device" biosver)

    $LOGGER_ECHO "$BASENAME" "The version of bootloader being updated: $source_binary_biosver"
    $LOGGER_ECHO "$BASENAME" "The version of current bootloader: $target_device_biosver"

    if [ "$source_binary_biosver" = "$target_device_biosver" ]; then
        $LOGGER_ECHO "$BASENAME" "Your bootloader version is the same as the version of bootloader being updated."
    fi
    _question "$assume_yes" "Do you want to continue? (y/N)"

    systemfailback=$(mx-system-mgmt -S state -V)
    if [ "$systemfailback" -eq "0" ]; then
        $LOGGER_ECHO "$BASENAME" "Bootloader update failure could result in device malfunction."
        $LOGGER_ECHO "$BASENAME" "It is recommended to enable system failback before the update."
        _question "$assume_yes" "Would you still like to continue without failback enabled? (y/N)"
    fi

    $LOGGER_ECHO "$BASENAME" "Start to upgrade bootloader..."
    if [ "$IOTHINX_FLAG" -eq "1" ]; then
        echo 0 >/sys/block/mmcblk2boot0/force_ro
        if ! dd if="$source_binary" of="$target_device" bs=1048576 seek=1 count=1; then
            $LOGGER_ECHO "$BASENAME" "Upgrade failed!"
            echo 1 >/sys/block/mmcblk2boot0/force_ro
            exit 1
        fi
        echo 1 >/sys/block/mmcblk2boot0/force_ro
    else
        if ! flashcp "$source_binary" "$target_device"; then
            $LOGGER_ECHO "$BASENAME" "Upgrade failed!"
            exit 1
        fi
    fi

    if ! verify_checksum "$target_device"; then
        $LOGGER_ECHO "$BASENAME" "Upgrade failed!"
        exit 1
    fi

    if ! systemctl start mx-bootloader-info-init; then
        $LOGGER_ECHO "$BASENAME" "Regen bootloader info failed!"
    fi

    $LOGGER_ECHO "$BASENAME" "Upgrade $target_device bootloader to version $source_binary_biosver successfully"
}

info() {
    local target=$1

    local compatible_model
    compatible_model="$(_query_info "$target" modelname)"
    local bootloader_version
    bootloader_version="$(_query_info "$target" biosver)"
    local sha256sum
    sha256sum="$(_query_info "$target" sha256sum)"
    local md5sum
    md5sum="$(_query_info "$target" md5sum)"

    if [ "${MBM_OPTION_JSON}" = "y" ]; then
        jq -n --arg compatible_model "${compatible_model}" --arg bootloader_version "${bootloader_version}" --arg sha256sum "${sha256sum}" --arg md5sum "${md5sum}" '{compatibleModel: $compatible_model, bootloaderVersion: $bootloader_version, sha256Sum: $sha256sum, md5Sum: $md5sum}'
    fi

    $LOGGER_ECHO "$BASENAME" "compatible model: $(_query_info "$target" modelname)"
    $LOGGER_ECHO "$BASENAME" "bootloader version: $(_query_info "$target" biosver)"
    $LOGGER_ECHO "$BASENAME" "sha256sum: $(_query_info "$target" sha256sum)"
    $LOGGER_ECHO "$BASENAME" "md5sum: $(_query_info "$target" md5sum)"
}

list() {
    json_str="["
    while read -r ub_file; do
        # Find all the "u-boot.bin" in STANDARD_UBOOT_DIR (/lib/firmware/moxa/bootloader) recursively.
        local compatible_model
        compatible_model="$(_query_info "$ub_file" modelname)"
        local bootloader_version
        bootloader_version="$(_query_info "$ub_file" biosver)"
        local sha256sum
        sha256sum="$(_query_info "$ub_file" sha256sum)"
        local md5sum
        md5sum="$(_query_info "$ub_file" md5sum)"

        # If information of the bootloader can't be obtained, according to _query_info(),
        # "Can't query XXX information" will be returned
        if [[ "$bootloader_version" =~ ^Can\'t || -z "$bootloader_version" ]]; then
            continue
        fi

        if [ "${MBM_OPTION_JSON}" = "y" ]; then
            json_str+="{\"bootloaderFilePath\": \"$ub_file\", \"compatibleModel\": \"$compatible_model\", \"bootloaderVersion\": \"$bootloader_version\", \"sha256sum\": \"$sha256sum\", \"md5sum\": \"$md5sum\"}"
        fi

        $LOGGER_ECHO "$BASENAME" "bootloader file path: ${ub_file}"
        $LOGGER_ECHO "$BASENAME" "compatible model: ${compatible_model}"
        $LOGGER_ECHO "$BASENAME" "bootloader version: ${bootloader_version}"
        $LOGGER_ECHO "$BASENAME" "sha256sum: ${sha256sum}"
        $LOGGER_ECHO "$BASENAME" "md5sum: ${md5sum}"

    done < <(find "$STANDARD_UBOOT_DIR" -type f -name "u-boot.bin" 2>/dev/null)

    json_str=${json_str//\}\{/\},\{}

    if [ "${MBM_OPTION_JSON}" = "y" ]; then
        json_str+="]"
        jq . <<<"$json_str"
    fi
}

detail() {
    local version_found=false
    while read -r ub_file; do
        # Find all the "u-boot.bin" in STANDARD_UBOOT_DIR (/lib/firmware/moxa/bootloader) recursively.
        local compatible_model
        compatible_model="$(_query_info "$ub_file" modelname)"
        local bootloader_version
        bootloader_version="$(_query_info "$ub_file" biosver)"
        local sha256sum
        sha256sum="$(_query_info "$ub_file" sha256sum)"
        local md5sum
        md5sum="$(_query_info "$ub_file" md5sum)"

        if [ "$version_detail" != "$bootloader_version" ]; then
            continue
        fi

        if [ "${MBM_OPTION_JSON}" = "y" ]; then
            jq -n --arg ub_file "${ub_file}" --arg cm "${compatible_model}" --arg bv "${bootloader_version}" --arg sha "${sha256sum}" --arg md "${md5sum}" '{bootloaderFilePath: $ub_file, compatibleModel: $cm, bootloaderVersion: $bv, sha256Sum: $sha, md5Sum: $md}'
        fi

        $LOGGER_ECHO "$BASENAME" "bootloader file path: ${ub_file}"
        $LOGGER_ECHO "$BASENAME" "compatible model: ${compatible_model}"
        $LOGGER_ECHO "$BASENAME" "bootloader version: ${bootloader_version}"
        $LOGGER_ECHO "$BASENAME" "sha256sum: ${sha256sum}"
        $LOGGER_ECHO "$BASENAME" "md5sum: ${md5sum}"

        version_found=true
        break

    done < <(find "$STANDARD_UBOOT_DIR" -type f -name "u-boot.bin" 2>/dev/null)

    if [ "$version_found" = false ]; then
        $LOGGER_ECHO "$BASENAME" "Bootloader file with version $version_detail not found"
        exit 1
    fi
}

parsing_options() {
    if ! OPTS=$(getopt -o if:yhJld: --long info,file:,yes,help,version,json,list,detail: -n "$INTERFACE_NAME" -- "$@"); then
        $HELPER_MENU "wrong" "upgrade"
        exit "${?}"
    fi
    eval set -- "$OPTS"

    while [ -n "$1" ]; do
        case "$1" in
        -i | --info)
            action=info
            shift
            ;;
        -f | --file)
            file="$2"
            shift
            shift
            ;;
        -y | --yes)
            assume_yes=y
            shift
            ;;
        -l | --list)
            action=list
            shift
            ;;
        -d | --detail)
            action=detail
            version_detail="$2"
            shift 2
            ;;
        -h | --help)
            $HELPER_MENU "upgrade"
            exit "${?}"
            ;;
        -J | --json)
            export MBM_OPTION_JSON=y
            shift
            ;;
        --)
            shift
            break
            ;;
        *)
            shift
            break
            ;;
        esac
    done

    if [ "$action" = "info" ]; then
        if [ -z "$file" ]; then
            target="$DEFAULT_UBOOT_DEV"
            $LOGGER_ECHO "$BASENAME" "Current bootloader information:"
        else
            if [ ! -f "$file" ]; then
                echo "$file it not exist"
                return 1
            fi
            target="$file"
            $LOGGER_ECHO "$BASENAME" "Bootloader file '$target' information:"
        fi
    fi

    if [ -z "$action" ]; then
        action=upgrade
        if [ -z "$file" ]; then
            return 1
        elif [ ! -f "$file" ]; then
            echo "$file it not exist"
            return 1
        fi
        target="$file"
    fi

    return 0
}

main() {
    if ! parsing_options "$@"; then
        $HELPER_MENU "wrong" "upgrade"
        exit "${?}"
    fi

    case "$action" in
    info)
        info "$target"
        ;;
    upgrade)
        upgrade "$target"
        ;;
    detail)
        detail
        ;;
    list)
        list
        ;;
    *) ;;
    esac

    return 0
}

main "$@"
