#!/bin/bash


#### usage ####

function exit_usage() {
    echo "Usage: fwupdate versions [-j]                     (lists available versions, optionally outputting json)"
    echo "       fwupdate current                           (shows the current version)"
    echo "       fwupdate download <version|url|file>       (downloads a firmware OS image)"
    echo "       fwupdate extract                           (extracts the downloaded firmware archive)"
    echo "       fwupdate flashboot                         (flashes the boot partition from extracted image)"
    echo "       fwupdate flashreboot                       (prepares for reboot + root partition flash)"
    echo "       fwupdate status                            (shows the current firmware updating status; see below)"
    echo "       fwupdate upgrade <version|url|file>        (performs all the operations necessary for upgrading)"
    echo ""
    echo "Statuses:"
    echo "       idle"
    echo "       downloading [version]"
    echo "       downloaded [version]"
    echo "       extracting [version]"
    echo "       extracted [version]"
    echo "       flashing boot [version]"
    echo "       boot flashed [version]"
    echo "       rebooting [version]"
    echo ""

    exit 1
}

if [[ -z "$1" ]]; then
    exit_usage
fi


#### configuration ####

set -a
set -e

SYS_VERSION_FILE=/etc/version
SYS_BOARD_FILE=/etc/board
OS_CONF_FILE=/etc/init.d/os_conf

MIN_FREE_DISK=500  # MB
FW_DIR=/data/.fwupdate

FW_FILE_GZ=${FW_DIR}/firmware.img.gz
FW_FILE_XZ=${FW_DIR}/firmware.img.xz
FW_FILE_EXTR=${FW_DIR}/firmware.img

VER_FILE=${FW_DIR}/version

BOOT_INFO_FILE=${FW_DIR}/boot_info
ROOT_INFO_FILE=${FW_DIR}/root_info

DOWNLOAD_STARTED_FILE=${FW_DIR}/download_started
DOWNLOAD_DONE_FILE=${FW_DIR}/download_done
EXTRACT_STARTED_FILE=${FW_DIR}/extract_started
EXTRACT_DONE_FILE=${FW_DIR}/extract_done
FLASH_BOOT_STARTED_FILE=${FW_DIR}/flash_boot_started
FLASH_BOOT_DONE_FILE=${FW_DIR}/flash_boot_done
FLASH_REBOOT_STARTED_FILE=${FW_DIR}/flash_reboot_started
ERROR_FILE=${FW_DIR}/error

TMP_BOOT_DIR=/tmp/fwupdate_boot
TMP_ROOT_DIR=/tmp/fwupdate_root

CURL_LOG_FILE=${FW_DIR}/curl.log
CURL_PID_FILE=${FW_DIR}/curl.pid

GUNZIP_LOG_FILE=${FW_DIR}/gunzip.log
GUNZIP_PID_FILE=${FW_DIR}/gunzip.pid

XZCAT_LOG_FILE=${FW_DIR}/xzcat.log
XZCAT_PID_FILE=${FW_DIR}/xzcat.pid

DD_LOG_FILE=${FW_DIR}/dd.log
DD_PID_FILE=${FW_DIR}/dd.pid

BOOT_LOOP="/dev/loop3"
ROOT_LOOP="/dev/loop4"

source ${OS_CONF_FILE}
source ${SYS_VERSION_FILE}


#### cleanup on exit ####

function cleanup_on_exit() {
    set +e

    if [[ -f /sbin/reboot.bak ]]; then
        rm -f /sbin/reboot
        mv /sbin/reboot.bak /sbin/reboot
    fi

    umount ${TMP_BOOT_DIR} 2>/dev/null
    umount ${TMP_ROOT_DIR} 2>/dev/null
    losetup -d ${BOOT_LOOP} 2>/dev/null
    losetup -d ${ROOT_LOOP} 2>/dev/null

    mount -T /etc/fstab.disk -o ro /boot 2>/dev/null
}


#### disk & partition devices ####

BOOT_DEV=$(mount | grep /boot | cut -d ' ' -f 1)
if [[ -n "${BOOT_DEV}" ]]; then
    ROOT_DEV=${BOOT_DEV:0:-1}2
    DATA_DEV=${BOOT_DEV:0:-1}3
else  # boot partition not mounted, trying with data partition
    DATA_DEV=$(mount | grep /data | cut -d ' ' -f 1)
    BOOT_DEV=${DATA_DEV:0:-1}1
    ROOT_DEV=${BOOT_DEV:0:-1}2
fi

DISK_DEV=${BOOT_DEV}
if [[ "${ROOT_DEV}" =~ ^([/a-z0-9]+)(p[0-9])$ ]]; then  # e.g. /dev/mmcblk0p2
    DISK_DEV=${BASH_REMATCH[1]}
elif [[ "${ROOT_DEV}" =~ ^([/a-z0-9]+)([0-9])$ ]]; then  # e.g. /dev/sdc2
    DISK_DEV=${BASH_REMATCH[1]}
else
    echo "cannot identify disk device from ${ROOT_DEV}" 1>&2
    exit 1
fi

function reallocate_boot_part() {
    current_boot_info=$(fdisk --bytes -l -o device,start,end,size ${DISK_DEV} | grep "${BOOT_DEV}")
    current_boot_info=(${current_boot_info})

    current_root_info=$(fdisk --bytes -l -o device,start,end,size ${DISK_DEV} | grep "${ROOT_DEV}")
    current_root_info=(${current_root_info})

    boot_info=($(cat ${BOOT_INFO_FILE}))

    if [[ ${current_boot_info[1]} == ${boot_info[0]} ]] &&
       [[ ${current_boot_info[2]} == ${boot_info[1]} ]]; then

        return  # all good
    fi

    echo "reallocating boot partition..."

    # check overlapping with root partition
    if [[ ${boot_info[1]} -ge ${current_root_info[1]} ]]; then
        echo "cannot reallocate boot partition: will overlap with root"
        return 1
    fi

    fdisk -w auto ${DISK_DEV} >/dev/null <<END
d
1
n
p
1
${boot_info[0]}
${boot_info[1]}
t
1
c
w
END
    sync
}

function reallocate_root_part() {
    current_root_info=$(fdisk --bytes -l -o device,start,end,size ${DISK_DEV} | grep "${ROOT_DEV}")
    current_root_info=(${current_root_info})

    current_data_info=$(fdisk --bytes -l -o device,start,end,size ${DISK_DEV} | grep "${DATA_DEV}")
    current_data_info=(${current_data_info})

    root_info=($(cat ${ROOT_INFO_FILE}))

    if [[ ${current_root_info[1]} == ${root_info[0]} ]] &&
       [[ ${current_root_info[2]} == ${root_info[1]} ]]; then

        return  # all good
    fi

    echo "reallocating root partition..."

    # check overlapping with data partition
    if [[ ${root_info[1]} -ge ${current_data_info[1]} ]]; then
        echo "cannot reallocate root partition: will overlap with data"
        return 1
    fi

    fdisk -w auto ${DISK_DEV} >/dev/null <<END
d
2
n
p
2
${root_info[0]}
${root_info[1]}
t
2
83
w
END
    sync
}


#### versions ####

function show_versions() {
    board=$(cat ${SYS_BOARD_FILE})
    show_json=$1

    # the /usr/libexec/list-versions-* helpers return a table with the following format:
    #     <version>|<prerelease>|<board>|<url>|<date|published>
    versions=$(FW_USERNAME=${OS_FIRMWARE_USERNAME} FW_PASSWORD=${OS_FIRMWARE_PASSWORD} \
               /usr/libexec/list-versions-${OS_FIRMWARE_METHOD} ${OS_FIRMWARE_REPO})

    for version in ${versions[@]}; do
        OIFS=${IFS}
        IFS="|"
        varr=(${version})
        IFS=${OIFS}
        if [[ "${OS_PRERELEASES}" == "false" && "${varr[1]}" == "true" ]]; then
            continue  # skip prereleases
        fi
        if [[ "${varr[5]}" == "false" ]]; then
            continue  # skip unpublished
        fi
        if [[ "${board}" != "${varr[2]}" ]]; then
            continue  # skip other boards
        fi

        if [[ "${show_json}" == "true" ]]; then
            echo "{\"version\": \"${varr[0]}\"," \
                  "\"url\": \"${varr[3]}\"," \
                  "\"prerelease\": ${varr[1]:-false}," \
                  "\"date\": \"${varr[4]}\"}"
        else
            echo ${varr[0]}
        fi
    done
}

function show_current() {
    echo ${OS_VERSION}
}


#### download ####

function do_download() {
    rm -rf ${FW_DIR}/*
    mkdir -p ${FW_DIR}

    rm -f ${DOWNLOAD_DONE_FILE}
    rm -f ${EXTRACT_STARTED_FILE}
    rm -f ${EXTRACT_DONE_FILE}
    rm -f ${FLASH_BOOT_STARTED_FILE}
    rm -f ${FLASH_BOOT_DONE_FILE}
    rm -f ${FLASH_REBOOT_STARTED_FILE}
    rm -f ${ERROR_FILE}
    touch ${DOWNLOAD_STARTED_FILE}

    echo "downloading..."

    # Look for local file first
    if [[ -f "$1" ]]; then
        version="custom"
        FNAME=`basename $1`
        FILEEXT=${FNAME##*.}
        DST_FNAME=""
        if [[ "${FILEEXT}" == "xz" ]]; then
            DST_FNAME=${FW_FILE_XZ}
        elif [[ "${FILEEXT}" == "gz" ]]; then
            DST_FNAME=${FW_FILE_GZ}
        fi
        if [[ -n "${DST_FNAME}" ]]; then
            cp -f $1 ${DST_FNAME}
            echo ${version} > ${VER_FILE}
            touch ${DOWNLOAD_DONE_FILE}
            return
        fi
    fi

    board=$(cat ${SYS_BOARD_FILE})
    url=$1
    version=$1

    if ! [[ "${url}" == http* ]]; then  # a version was given
        url=$(show_versions true | jq -r ". | select(.version==\"${version}\") | .url")
    else
        version="custom"
    fi

    if [[ -z "${url}" ]]; then
        echo "no such version" | tee 1>&2 ${ERROR_FILE}
        rm ${DOWNLOAD_STARTED_FILE}
        exit 1
    fi

    free_disk=$(df /data | tail -n 1 | tr -s ' ' | cut -d ' ' -f 4)
    if [[ "${free_disk}" -lt $((MIN_FREE_DISK * 1024)) ]]; then
        echo "not enough disk space" | tee 1>&2 ${ERROR_FILE}
        rm ${DOWNLOAD_STARTED_FILE}
        exit 1
    fi

    outfile=${FW_FILE_GZ}
    format=$(echo ${url} | sed -rn 's/.*\.img\.([a-z]+)$/\1/ p')
    if [[ "${format}" == "xz" ]]; then
        outfile=${FW_FILE_XZ}
    fi

    echo ${version} > ${VER_FILE}

    curl_opts="-S -f -L"
    if [[ -n "${OS_FIRMWARE_USERNAME}" ]]; then
        curl_opts+=" --user ${OS_FIRMWARE_USERNAME}:${OS_FIRMWARE_PASSWORD}"
    fi

    curl ${curl_opts} -o ${outfile} "${url}" &> ${CURL_LOG_FILE} &
    pid=$!
    echo ${pid} > ${CURL_PID_FILE}

    set +e  # don't exit on error just yet
    wait ${pid}

    if [[ "$?" != 0 ]]; then
        echo "download failed" | tee 1>&2 ${ERROR_FILE}
        rm ${DOWNLOAD_STARTED_FILE}
        exit 1
    fi

    touch ${DOWNLOAD_DONE_FILE}
    set -e
}


#### extract ####

function run_pre_upgrade() {
    which losetup &>/dev/null || return 0

    boot_info=($(cat ${BOOT_INFO_FILE}))
    root_info=($(cat ${ROOT_INFO_FILE}))
    pre_upgrade="${TMP_ROOT_DIR}/usr/share/pre-upgrade/*"

    mkdir -p ${TMP_BOOT_DIR}
    mkdir -p ${TMP_ROOT_DIR}
    losetup -o $((boot_info[0] * 512)) ${BOOT_LOOP} ${FW_FILE_EXTR}
    losetup -o $((root_info[0] * 512)) ${ROOT_LOOP} ${FW_FILE_EXTR}
    mount ${BOOT_LOOP} ${TMP_BOOT_DIR}
    mount ${ROOT_LOOP} ${TMP_ROOT_DIR}

    if [[ -d ${pre_upgrade} ]]; then
        for script in ${pre_upgrade}/*.sh; do
            echo "running pre-upgrade script $(basename ${script})"
            if [[ -x ${script} ]] && ! ${script}; then
                # non-zero exit status of pre-upgrade script indicates that
                # the upgrade process must not be continued

                echo "aborted by pre-upgrade script" | tee 1>&2 ${ERROR_FILE}
                return 1
            fi
        done
    fi

    umount ${TMP_BOOT_DIR}
    umount ${TMP_ROOT_DIR}
    losetup -d ${BOOT_LOOP}
    losetup -d ${ROOT_LOOP}
}

function do_extract() {
    if ! [[ -f ${FW_FILE_GZ} || -f ${FW_FILE_XZ} ]] || ! [[ -f ${DOWNLOAD_DONE_FILE} ]]; then
        echo "firmware file not downloaded" | tee 1>&2 ${ERROR_FILE}
        exit 1
    fi

    rm -f ${EXTRACT_DONE_FILE}
    rm -f ${FLASH_BOOT_STARTED_FILE}
    rm -f ${FLASH_BOOT_DONE_FILE}
    rm -f ${FLASH_REBOOT_STARTED_FILE}
    rm -f ${ERROR_FILE}
    rm -f ${FW_FILE_EXTR}
    touch ${EXTRACT_STARTED_FILE}

    echo "extracting..."

    trap cleanup_on_exit EXIT

    format="gz"
    if [[ -f ${FW_FILE_XZ} ]]; then
        format="xz"
    fi

    rm -f ${FW_FILE_EXTR}
    rm -f ${GUNZIP_PID_FILE} ${XZCAT_PID_FILE}

    if [[ "${format}" == "xz" ]]; then
        DECOMPRESS_LOG_FILE=${XZCAT_LOG_FILE}
        DECOMPRESS_PID_FILE=${XZCAT_PID_FILE}
        xzcat ${FW_FILE_XZ} > ${FW_FILE_EXTR} 2>${XZCAT_LOG_FILE} &
    elif [[ "${format}" == "gz" ]]; then
        DECOMPRESS_LOG_FILE=${GUNZIP_LOG_FILE}
        DECOMPRESS_PID_FILE=${GUNZIP_PID_FILE}
        gunzip -k -c ${FW_FILE_GZ} > ${FW_FILE_EXTR} 2>${GUNZIP_LOG_FILE} &
    else
        echo "firmware compression format ${format} not supported" 1>&2 | tee 1>&2 ${ERROR_FILE}
        exit 1
    fi

    pid=$!
    echo ${pid} > ${DECOMPRESS_PID_FILE}
    wait ${pid}

    if [[ "$?" != 0 ]]; then
        cat ${DECOMPRESS_LOG_FILE}
        exit 1
    fi

    # TODO verify hash

    boot_info=$(fdisk --bytes -l -o device,start,end,size ${FW_FILE_EXTR} | grep "${FW_FILE_EXTR}1")
    boot_info=(${boot_info})
    boot_start=${boot_info[1]}
    boot_end=${boot_info[2]}
    boot_size=$((boot_info[3] / 512))

    root_info=$(fdisk --bytes -l -o device,start,end,size ${FW_FILE_EXTR} | grep "${FW_FILE_EXTR}2")
    root_info=(${root_info})
    root_start=${root_info[1]}
    root_end=${root_info[2]}
    root_size=$((root_info[3] / 512))

    echo ${boot_start} ${boot_end} ${boot_size} > ${BOOT_INFO_FILE}
    echo ${root_start} ${root_end} ${root_size} > ${ROOT_INFO_FILE}

    run_pre_upgrade

    touch ${EXTRACT_DONE_FILE}
}


#### flash boot ####

function do_flash_boot() {
    if ! [[ -f ${FW_FILE_EXTR} ]] || ! [[ -f ${EXTRACT_DONE_FILE} ]]; then
        echo "extracted firmware not present" | tee 1>&2 ${ERROR_FILE}
        return 1
    fi

    rm -f ${FLASH_BOOT_DONE_FILE}
    rm -f ${FLASH_REBOOT_STARTED_FILE}
    rm -f ${ERROR_FILE}
    touch ${FLASH_BOOT_STARTED_FILE}

    echo "flashing boot..."

    trap cleanup_on_exit EXIT

    board=$(cat ${SYS_BOARD_FILE})
    boot_info=($(cat ${BOOT_INFO_FILE}))

    cp -r /boot ${FW_DIR}/old_boot
    umount /boot

    # prevent unwanted reboots during upgrade
    mount -o remount,rw /
    mv /sbin/reboot /sbin/reboot.bak
    ln -s /bin/true /sbin/reboot

    reallocate_boot_part

    dd if=${FW_FILE_EXTR} \
       skip=$((boot_info[0] / 2048)) \
       of=${BOOT_DEV} bs=1048576 \
       count=$((boot_info[2] / 2048)) &>${DD_LOG_FILE} &
    pid=$!
    echo ${pid} > ${DD_PID_FILE}
    wait ${pid}

    mount -T /etc/fstab.disk -o rw /boot

    # the /usr/libexec/fw-restore-boot-cfg script, if present, takes the old (backup) boot dir as argument
    # and should restore any /boot configuration that needs to be preserved across updates
    # from the old boot dir to the current (new) /boot dir
    if [[ -x /usr/libexec/fw-restore-boot-cfg ]]; then
        /usr/libexec/fw-restore-boot-cfg ${FW_DIR}/old_boot 2>/dev/null || true
    fi

    touch ${FLASH_BOOT_DONE_FILE}
}


#### flash reboot ####

function do_flash_reboot() {
    if ! [[ -f ${ROOT_INFO_FILE} ]] || ! [[ -f ${FLASH_BOOT_DONE_FILE} ]]; then
        echo "root partition info file not present" | tee 1>&2 ${ERROR_FILE}
        return 1
    fi

    rm -f ${ERROR_FILE}
    touch ${FLASH_REBOOT_STARTED_FILE}

    echo "preparing for reboot..."

    trap cleanup_on_exit EXIT

    reallocate_root_part

    board=$(cat ${SYS_BOARD_FILE})

    # the /usr/libexec/fw-prepare-boot script should be present and should
    # make the necessary changes to the current boot configuration so that
    # the firmware update initramfs will be used by the next boot
    mount -o remount,rw /boot
    if [[ -x /usr/libexec/fw-prepare-boot ]]; then
        /usr/libexec/fw-prepare-boot
    fi

    sync
    busybox reboot &
    sleep 10
    echo b > /proc/sysrq-trigger

    exit 0
}


#### status ####

function show_status() {
    if [[ -f ${VER_FILE} ]]; then
        new_version=$(cat ${VER_FILE})
    fi

    if [[ -f ${ERROR_FILE} ]]; then
        echo -n "error: "
        cat ${ERROR_FILE}
    elif [[ -f ${FLASH_REBOOT_STARTED_FILE} ]]; then
        echo "rebooting [${new_version}]"
    elif [[ -f ${FLASH_BOOT_DONE_FILE} ]]; then
        echo "boot flashed [${new_version}]"
    elif [[ -f ${FLASH_BOOT_STARTED_FILE} ]]; then
        echo "flashing boot [${new_version}]"
    elif [[ -f ${EXTRACT_DONE_FILE} ]]; then
        echo "extracted [${new_version}]"
    elif [[ -f ${EXTRACT_STARTED_FILE} ]]; then
        echo "extracting [${new_version}]"
    elif [[ -f ${DOWNLOAD_DONE_FILE} ]]; then
        echo "downloaded [${new_version}]"
    elif [[ -f ${DOWNLOAD_STARTED_FILE} ]]; then
        echo "downloading"
    else
        echo "idle"
    fi
}


#### upgrade ####

function do_upgrade() {
    echo "upgrading to $1"

    do_download "$1"
    show_status

    do_extract
    show_status

    do_flash_boot
    show_status

    do_flash_reboot
}


#### main ####

case "$1" in
    versions)
        show_json="false"
        test "$2" == "-j" && show_json="true"
        show_versions ${show_json}
        ;;

    current)
        show_current
        ;;

    download)
        if [[ -z "$2" ]]; then
            exit_usage
        fi

        do_download "$2"
        show_status
        ;;

    extract)
        do_extract
        show_status
        ;;

    flashboot)
        do_flash_boot
        show_status
        ;;

    flashreboot)
        do_flash_reboot
        ;;

    status)
        show_status
        ;;

    upgrade)
        if [[ -z "$2" ]]; then
            exit_usage
        fi

        do_upgrade "$2"
        ;;

    *)
        exit_usage
        ;;
esac

