#!/bin/bash
#
# This script is used to upgrade a bank using the contents of a tarball
#

# globals
debug=0       # enabled by --debug or --dry-run
dry_run=0     # enabled by --dry-run
force=0       # enabled by --force

# exit codes (constants - don't change these, as they will be used in app code)
EXIT_SUCCESS=0
EXIT_INVALID_ARGS=1
EXIT_NO_SOURCE=2
EXIT_NOT_ROOT=3
EXIT_COMMAND_FAILED=4

# -----------------------------------------------------------------------------
print_usage_and_exit() {
    echo "Usage: $this_script [option] <A|B> <tarball>"
    echo
    echo -e "\tThis script will upgrade the given bank with the contents of the given tarball."
    echo
    echo "Options:"
    echo -e "\t--debug    displays debug information while running"
    echo -e "\t--dry-run  shows commands that will be run (but does nothing)"
    echo -e "\t--force    erases and reformats the target bank before starting"
    echo -e "\t--help     shows this message"
    echo
    exit $EXIT_INVALID_ARGS
}

# -----------------------------------------------------------------------------
command_wrapper() {
    local comment=$1; shift
    local command_line="$*"

    if [ "$debug" -eq "1" ] ; then
        echo "$comment"
        echo "$command_line"
        echo
    fi

    if [ "$dry_run" -eq "1" ] ; then
        return 0
    fi

    if [ "$debug" -eq "1" ] ; then
        $command_line
    else
        $command_line &> /dev/null
    fi

    return "$?"
}

# -----------------------------------------------------------------------------
exit_on_error() {
    local comment=$1

    command_wrapper "$@"

    if [ "$?" -ne "0" ] ; then
        echo "ERROR: $comment failed"

        # try to mount things the way they are supposed to be mounted
        command_wrapper "mount app ro"    mount $dst_app_mount -o remount,ro

        if copy_kernel; then
            command_wrapper "mount kernel ro" mount $dst_kern_mount -o remount,ro
        fi

        exit $EXIT_COMMAND_FAILED
    fi
}

# -----------------------------------------------------------------------------
extract_lzma_tarball() {
    cat $1 | unlzma -c | tar -xpf -
}

# ------------------------------------------------------------------------------
spiboot_supported() {
    spiboot_info='/sys/bus/platform/devices/board_info/spiboot'
    [ -r "${spiboot_info}" ] && [ "$(<${spiboot_info})" == "1" ]
}

# ------------------------------------------------------------------------------
copy_kernel() {
    # For boards that boot from SPI flash the kernel isn't in a separate partition.
    ! spiboot_supported
}

#--------------------------------------------------------------------
find_mtd_device() {
    for x in "$@"; do
        for y in /sys/class/mtd/mtd*; do
            if [ -f $y/name ] && [ "$(<$y/name)" == "$x" ]; then
                basename $y
                return 0
            fi
        done
    done

    return 1
}

# -----------------------------------------------------------------------------
flush_cache() {
    sync
    echo 3 > /proc/sys/vm/drop_caches
}

# -----------------------------------------------------------------------------
md5sum_check() {
    # run the md5sum check on the files in the manifest
    md5sum -s -c ./MANIFEST
}

# -----------------------------------------------------------------------------
get_part() {
    # Convert partitions of the format /dev/mtdXblock or mtd:name to their
    # character device counterparts /dev/mtdX or /dev/mtd/name.

    local mnt=$1
    local part=${mnt/block}

    if [ "$mnt" != "$part" ]; then
        echo $part
    else
        echo /dev/mtd/${mnt/#mtd:}
    fi
}

# -----------------------------------------------------------------------------
# Main script

# Take note of our name as invoked
this_script=$(basename $0)

# handle options
while :; do
    case "$1" in
        "--help"|"-h"|"-?"|"help") print_usage_and_exit ;;
        "--debug") debug=1 ; shift ;;
        "--dry-run") debug=1 ; dry_run=1; shift ;;
        "--force") force=1; shift ;;
        *) break ;;
    esac
done

if [ "$#" -ne "2" ] ; then
    echo "ERROR: wrong number of parameters"
    print_usage_and_exit
fi

active_image="$1"
src_image="$2"

if [ "$(whoami)" != "root" ] ; then
    echo "ERROR: must run $this_script as root"
    exit $EXIT_NOT_ROOT
fi

case "$active_image" in
    a|A) active_image="A"
         dst_app_mount=/mnt/imageA
         dst_kern_mount=/mnt/kernelA
         src_kernel=${dst_app_mount}/boot/*
         dst_kernel=${dst_kern_mount}/
         src_uboot=${dst_app_mount}/ciena/images/u-boot-A.img

         dst_uboot=/dev/$(find_mtd_device u-boot-a "Uboot A")
         if [ "$?" -ne "0" ] ; then
             echo "Unable to locate mtd device for u-boot a"
             exit $EXIT_COMMAND_FAILED
         fi
         ;;
    b|B) active_image="B"
         dst_app_mount=/mnt/imageB
         dst_kern_mount=/mnt/kernelB
         src_kernel=${dst_app_mount}/boot/*
         dst_kernel=${dst_kern_mount}/
         src_uboot=${dst_app_mount}/ciena/images/u-boot-B.img

         dst_uboot=/dev/$(find_mtd_device u-boot-b "Uboot B")
         if [ $? -ne 0 ]; then
             echo "Unable to locate mtd device for u-boot b"
             exit $EXIT_COMMAND_FAILED
         fi
         ;;
    *)   active_image=""
         ;;
esac

# Boards that boot from SPI flash use a special U-Boot image.
if spiboot_supported; then
    src_uboot=${dst_app_mount}/ciena/images/u-boot-spi.img
fi

if [ -z "$active_image" ] ; then
    echo "ERROR: unable to determine source bank"
    exit $EXIT_NO_SOURCE
fi

if [ ! -f "$src_image" ] ; then
    echo "ERROR: $src_image missing"
    exit $EXIT_NO_SOURCE
fi

case "$src_image" in
    *.gz)
         extract_cmd="tar -xzpf $src_image"
         ;;
    *.xz)
         extract_cmd="extract_lzma_tarball $src_image"
         ;;
    *)
         echo "ERROR: $src_image has invalid extension"
         exit $EXIT_NO_SOURCE
         ;;
esac

if [ "$debug" -eq "1" ] ; then
    echo "active_image  = $active_image"
    echo "dst_app_mount = $dst_app_mount"
    echo "dst_kern_mount= $dst_kern_mount"
    echo "src_kernel    = $src_kernel"
    echo "dst_kernel    = $dst_kernel"
    echo "src_uboot     = $src_uboot"
    echo "dst_uboot     = $dst_uboot"
    echo "extract_cmd   = $extract_cmd"
fi

# If the force option is specified, then unmount the destination partitions
# and erase them.  The -j option to flash_eraseall will format the partition
# for use with jffs2.
#
if copy_kernel && [ "$force" -eq "1" ] ; then

    # resolve the flash partition behind the fs mount
    dst_app_part=
    dst_kern_part=
    while read -a __fs_entry; do
        if [ "${__fs_entry[2]}" = "jffs2" ]; then
            case ${__fs_entry[1]} in
                ${dst_app_mount})  dst_app_part=$(get_part ${__fs_entry[0]}) ;;
                ${dst_kern_mount}) dst_kern_part=$(get_part ${__fs_entry[0]}) ;;
                *) ;;
            esac
        fi
    done < /proc/mounts

    if [ "$debug" -eq "1" ]; then
        echo "dst_kern_part = $dst_kern_part"
        echo "dst_app_part  = $dst_app_part"
    fi

    if [ -n "${dst_kern_part}" ]; then
        command_wrapper "unmount kernel partition" umount $dst_kern_mount
        command_wrapper "erase kernel partition" nice flash_eraseall -qj $dst_kern_part
        rc="$?"
        exit_on_error   "mount kernel partition" mount $dst_kern_mount
        if [ "$rc" -ne "0" ] ; then
            echo "failed to erase kernel partition"
            exit $EXIT_COMMAND_FAILED
        fi
    fi

    if [ -n "${dst_app_part}" ]; then
        command_wrapper "unmount app partition" umount $dst_app_mount
        command_wrapper "erase app partition" nice flash_eraseall -qj $dst_app_part
        rc="$?"
        exit_on_error   "mount app partition"   mount $dst_app_mount
        if [ "$rc" -ne "0" ] ; then
            echo "failed to erase app partition"
            exit $EXIT_COMMAND_FAILED
        fi
    fi
fi

# Main protect logic
#
# Ensure that the u-boot partition is erased first, as this marks the partition
# as bad, in case we get interrupted or encounter an error.
#
exit_on_error "erase u-boot partition" flash_eraseall -q $dst_uboot
exit_on_error "remount app rw"         mount $dst_app_mount -o remount,rw
exit_on_error "cd to $dst_app_mount"   cd $dst_app_mount
exit_on_error "remove old files"       rm -rf ${dst_app_mount}/*
exit_on_error "extract new files"      $extract_cmd
# flush the cache to ensure the md5sum runs against the flash content
exit_on_error "flushing cache"         flush_cache
exit_on_error "remount app ro"         mount $dst_app_mount -o remount,ro
exit_on_error "md5sum check"           md5sum_check

if copy_kernel; then
    exit_on_error "remount kernel rw"      mount $dst_kern_mount -o remount,rw
    exit_on_error "remove old kernel"      rm -rf ${dst_kern_mount}/*
    exit_on_error "copy kernel"            cp $src_kernel $dst_kernel
    # flush the cache again to ensure the diff runs against the flash content
    exit_on_error "flushing cache"         flush_cache
    exit_on_error "remount kernel ro"      mount $dst_kern_mount -o remount,ro
    exit_on_error "checking kernel"        diff -r ${src_kernel%\/*} $dst_kernel
fi

exit_on_error "copy u-boot"            dd bs=128k if=$src_uboot of=$dst_uboot
exit_on_error "syncing app"            sync

exit $EXIT_SUCCESS
