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

# Do not make . files invisible. This allows rm -rf * to really clean up.
shopt -s dotglob

# 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

# vm dirty byte counts
vm_dirty_bytes_file="/proc/sys/vm/dirty_bytes"
vm_dirty_ratio_file="/proc/sys/vm/dirty_ratio"
vm_dirty_bytes_val=0
vm_dirty_ratio_val=0
vm_dirty_bytes_min=8192

# -----------------------------------------------------------------------------
print_usage_and_exit() {
    echo "Usage: $this_script [option] <A|B> <tarball> [optional 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--import   run \"source ${BASH_SOURCE##*/} --import\" to import functions"
    echo -e "\t--help     shows this message"
    echo
    exit $EXIT_INVALID_ARGS
}

# -----------------------------------------------------------------------------
get_board_name_caliondo_5()
{
    case "$(cat /sys/bus/platform/devices/board_info/pwe)" in

        0)
            echo '3930ET' ;;
        1)
            echo '3932' ;;
        *)
            echo 'unknown' ;;
    esac
}

# -----------------------------------------------------------------------------
get_board_name_caliondo()
{
    case "$(cat /sys/bus/platform/devices/board_info/id)" in

        1)
            echo '5150' ;;
        2)
            echo '3916' ;;
        3)
            echo '3930' ;;
        4)
            echo '3931' ;;
        5)
            get_board_name_caliondo_5 ;;
        6)
            echo '5142' ;;
        7)
            echo '5160' ;;
        8)
            echo '3942' ;;
        9)
            echo '3938' ;;
        *)
            echo 'unknown' ;;
    esac
}

# -----------------------------------------------------------------------------
get_board_name_acorn()
{
    local model_file="/proc/device-tree/model"

    if [ -f "$model_file" ] ; then
        cat $model_file | tr -d '\0'
    else
        echo 'unknown'
    fi
}

# -----------------------------------------------------------------------------
get_board_name()
{
    case "$(cat /family)" in
        "caliondo") get_board_name_caliondo ;;
        "dernhelm"| \
        "eredan")   get_board_name_acorn ;;
        *) echo 'unknown' ;;
    esac
}

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

    # CONSOLE_COLLECTOR must match definition in rcS
    local CONSOLE_COLLECTOR=/tmp/log/console_collector

    # if the collector isn't a pipe then don't use it
    if [ ! -p "$CONSOLE_COLLECTOR" ] ; then
        CONSOLE_COLLECTOR=/dev/null
    fi

    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 2> $CONSOLE_COLLECTOR
    fi

    return "$?"
}

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

    if ! command_wrapper "$@" ; 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

        vm_dirty_restore

        exit $EXIT_COMMAND_FAILED
    fi
}

# -----------------------------------------------------------------------------
die() {
    echo $2
    exit $1
}

# -----------------------------------------------------------------------------
vm_dirty_set() {
    echo $2 > $1
}

# -----------------------------------------------------------------------------
vm_dirty_shrink() {

    if [ "${vm_dirty_bytes_val}" != "0" ] || \
       [ "${vm_dirty_ratio_val}" != "0" ]; then
        command_wrapper "setting ${vm_dirty_bytes_file} to ${vm_dirty_bytes_min}" \
            vm_dirty_set ${vm_dirty_bytes_file} ${vm_dirty_bytes_min}
    fi
}

# -----------------------------------------------------------------------------
vm_dirty_restore() {
    # either dirty_bytes or dirty_ratio must be set
    # the other one must be zero
    if [ "${vm_dirty_bytes_val}" != "0" ]; then
        command_wrapper "restoring ${vm_dirty_bytes_file} to ${vm_dirty_bytes_val}" \
            vm_dirty_set ${vm_dirty_bytes_file} ${vm_dirty_bytes_val}
    elif [ "${vm_dirty_ratio_val}" != "0" ]; then
        command_wrapper "restoring ${vm_dirty_ratio_file} to ${vm_dirty_ratio_val}" \
            vm_dirty_set ${vm_dirty_ratio_file} ${vm_dirty_ratio_val}
    fi
}

# -----------------------------------------------------------------------------
extract_lzma_tarball() {
    local image=$1; shift
    xz -d -c < "$image" | tar -xpf - "$@"
}

# -----------------------------------------------------------------------------
list_lzma_tarball() {
    local image=$1; shift
    xz -d -c < "$image" | tar -tf - "$@"
}

# ------------------------------------------------------------------------------
spiboot_supported() {

    # All current dernhelm platforms should be treated as spi boot
    if [ "$(cat /family)" == "dernhelm" ] ; then
        return 0
    fi

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

# ------------------------------------------------------------------------------
is_5150()
{
    local board_id='/sys/bus/platform/devices/board_info/id'
    [ -r "${board_id}" ] && [ "$(<${board_id})" == "1" ]
}

# ------------------------------------------------------------------------------
copy_kernel() {
    # This function must agree with all the copy_kernel() in the
    # target-specific board_lib.sh scripts.

    # For boards that boot from SPI flash the kernel isn't in a
    # separate partition.
    spiboot_supported && return 1

    # U-Boot takes the kernel from the main fs image on caliondo NAND
    # boards, and the 5150 is the only caliando board with a NOR
    # flash.
    [[ "$(cat /family)" == "caliondo" ]] && ! is_5150 && return 1

    # Everything else (brego and 5150): copy the kernel.
    return 0
}

#--------------------------------------------------------------------
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
}

# -----------------------------------------------------------------------------
boot_bin_version() {
    local _filename="$1"
    local _ver
    # Save the current pipefail setting
    local _pipefail="$(set +o | grep pipefail)"

    set -o pipefail
    _ver="$(/ciena/bin/bootbinExtract "$_filename" | tr -d '\0' | head -n 1)" \
        || _ver="0"

    # Restore the pipefail setting
    eval ${_pipefail}

    echo "${_ver}"
}

# -----------------------------------------------------------------------------
update_bootloader() {
    local _boot_mount="/mnt/boot"
    grep " ${_boot_mount} " /proc/mounts \
        || command_wrapper "mount bootfs ro"  mount "${_boot_mount}" \
	    || return "$?"
    local _rc
    local _new_boot_bin="${1}/ciena/images/BOOT.BIN"
    local _cur_version="$(boot_bin_version "${_boot_mount}/BOOT.BIN")"
    local _new_version="$(boot_bin_version "${_new_boot_bin}")"

    command_wrapper "mount bootfs rw"         mount -o remount,rw "${_boot_mount}" \
        || return "$?"

    # Remove the uboot environment and uboot xtras files. 
    command_wrapper "remove saved u-boot env" rm -f "${_boot_mount}/uboot.env" &&\
    command_wrapper "remove u-boot xtras"     rm -f "${_boot_mount}/uenvxtra.txt" &&\
    command_wrapper "sync filesystems"        sync
    _rc="$?"
    if [ "$_rc" -ne "0" ] ; then
        echo "INFO: Unable to clean up e-boot environment"
        command_wrapper "umount bootfs"           mount -o remount,ro "${_boot_mount}"
        return $_rc
    fi

    if [ "${_cur_version}" -eq "0" ] ||\
       [ "${_new_version}" -eq "0" ]; then
        echo "INFO: Bad versions, existing version ${_cur_version}, new version ${_new_version}"
        command_wrapper "umount bootfs"           mount -o remount,ro "${_boot_mount}"
        return 1
    fi

    if [ "${_cur_version}" -eq "${_new_version}" ] ; then
        echo "INFO: No need to upgrade, version ${_cur_version} is already installed"
        command_wrapper "umount bootfs"           mount -o remount,ro "${_boot_mount}" \
            || return "$?"
        return 0
    fi

    if [ "${_cur_version}" -gt "${_new_version}" ] ; then
        echo "INFO: No need to upgrade, existing version ${_cur_version} is newer than ${_new_version}"
        command_wrapper "umount bootfs"           mount -o remount,ro "${_boot_mount}" \
            || return "$?"
        return 0
    fi

    command_wrapper "backup BOOT.BIN"         cp "${_boot_mount}/BOOT.BIN" "${_boot_mount}/BOOT0001.BIN" &&\
    command_wrapper "update BOOT.BIN"         cp "${_new_boot_bin}" "${_boot_mount}/BOOT.BIN" &&\
    command_wrapper "sync filesystems"        sync
    _rc="$?"

    command_wrapper "umount bootfs"           mount -o remount,ro "${_boot_mount}" \
        || return "$?"

    return "${_rc}"

}

# -----------------------------------------------------------------------------
sha256sum_check() {
    # See if busybox quiet works otherwise use coreutils option
    local quiet='-s'
    sha256sum ${quiet} -c < /dev/null &> /dev/null || quiet='--quiet'

    # run sha256sum check on the files in the manifest
    sha256sum ${quiet} -c ./MANIFEST_SHA256
}

# -----------------------------------------------------------------------------
get_mtd_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
}

# -----------------------------------------------------------------------------
# $1 -- directory $2- install command
install_board_specific_files() {
     local dir=$1
     local cmd=$2
     local tmpManifest
     tmpManifest=$(mktemp -t temp$$.XXXXXX)
     exit_on_error "remount $dir rw"        mount $dir -o remount,rw
     exit_on_error "cd to $dir"             cd $dir

     # save the MANIFEST file from the previous tarball
     exit_on_error "mv MANIFEST(sha256)"     mv $dir/MANIFEST_SHA256 $tmpManifest

     exit_on_error "extract new files"      $cmd
     exit_on_error "flushing cache"         flush_cache

     if [ -f "$dir/MANIFEST_SHA256" ]; then
        exit_on_error "sha256sum check"    sha256sum_check
        chmod 644 $dir/MANIFEST_SHA256
        cat $tmpManifest >> $dir/MANIFEST_SHA256
        rm  $tmpManifest
     else
        exit_on_error "restore manifest(sha256)"  mv $tmpManifest $dir/MANIFEST_SHA256
     fi
     exit_on_error "unmark partial manifest" manifest_remove_blackmark
     exit_on_error "flushing cache"         flush_cache
     exit_on_error "remount $dir ro"        mount $dir -o remount,ro
}

# -----------------------------------------------------------------------------
flash_format() {
    local jffs2=""

    case $1 in
        "--jffs2")  jffs2="--jffs2"; shift ;;
    esac

    local device=$1

    [ -c "$device" ] || die $EXIT_COMMAND_FAILED "No such device ($device)"

    # Support for --version was atomically introduced when mtd-utils
    # deprecated flash_eraseall.
    if flash_erase --version > /dev/null 2>&1; then
        nice flash_erase    --quiet $jffs2 $device 0 0
    else
        nice flash_eraseall --quiet $jffs2 $device
    fi

}

# -----------------------------------------------------------------------------
find_ubi_volume() {
    # $1: volume name
    for x in /sys/class/ubi/ubi*; do
        if [ -f $x/name ] && [ "$(<$x/name)" == "$1" ]; then
            basename $x
            return 0
        fi
    done

    return 1
}


# -----------------------------------------------------------------------------
ubifs_truncate() {
    local volume=$1
    local vid

    vid=$(find_ubi_volume $volume) \
        || die $EXIT_COMMAND_FAILED "Failed to find ubi volume $volume"

    nice ubiupdatevol /dev/${vid} --truncate

}

# -----------------------------------------------------------------------------
create_ramfs_manifest()
{
    local _rfs=${1}
    local _ramfs=${2}
    local _tmp_manifest
    local _f
    local _rc
    _tmp_manifest=$(mktemp -t temp$$.XXXXXX)

    # Grab the manifest entries from the rootfs and use them to
    # create a manifest for the ramfs mount.  Note that our
    # current rootfs manifest does not cover symbolic links.
    for _f in $(cd ${_ramfs} && find . -type f) ; do
        grep -E ${_f}$ ${_rfs}/MANIFEST_SHA256 >> ${_tmp_manifest}
        _rc="$?"
        if [ "${_rc}" != "0" ] ; then
            rm ${_tmp_manifest}
            echo "ERROR: failed to find ${_f} in ${_rfs}/MANIFEST_SHA256" >&2
            return "${_rc}"
        fi
    done
    mv ${_tmp_manifest} ${_ramfs}/MANIFEST_SHA256
}

# -----------------------------------------------------------------------------
ramfs_fixup()
{
    # U-Boot expects images in /dev/mmcblk0p[56]:
    #   /boot/FITImage.itb
    #   /ciena/fpga/*
    #
    # When /dev/mmcblk0p[56] is mounted as /mnt/ramfs[AB], the images
    # must be copied manually.
    local _rfs=${1}
    local _ramfs=${_rfs/image/ramfs}

    if [ ! -d ${_ramfs} ]; then
        return 0
    fi

    command_wrapper "mount ramfs rw"     mount ${_ramfs} -o rw \
        || return "$?"
    command_wrapper "remove old files"   rm -rf ${_ramfs}/* &&\
    command_wrapper "create dest dir"    mkdir -p ${_ramfs}/ciena &&\
    command_wrapper "copy FIT image"     cp -a ${_rfs}/boot ${_ramfs} &&\
    command_wrapper "copy fpga images"   cp -a ${_rfs}/ciena/fpga ${_ramfs}/ciena &&\
    command_wrapper "ramfs manifest"     create_ramfs_manifest ${_rfs} ${_ramfs}
    _rc="$?"

    command_wrapper "umount ramfs"       umount ${_ramfs} \
        || return "$?"

    return "${_rc}"

}

# -----------------------------------------------------------------------------
erase_partition() {
    local name=$1
    local mount=$2

    local fs_entry
    local fs_type
    local device
    local rc

    while read -a fs_entry; do
        if [ ${fs_entry[1]} = $mount ]; then
            device=${fs_entry[0]}
            fs_type=${fs_entry[2]}
            break
        fi
    done < /proc/mounts

    if [ "$debug" -eq "1" ] ; then
        echo "erase_partition()"
        echo "$name device  = $device"
        echo "$name fs_type = $fs_type"
    fi

    command_wrapper "unmount $name partition" umount $mount

    case "$fs_type" in
        jffs2) command_wrapper "format $name partition" \
                        flash_format --jffs2 $(get_mtd_part $device) ;;
        ubifs) command_wrapper "truncate $name partition" \
                        ubifs_truncate ${device#*:} ;;
        *)     die $EXIT_COMMAND_FAILED "Unknown filesystem for $name" ;;
    esac

    rc=$?

    # ubifs needs to be mounted read/write the first time
    exit_on_error "mount $name partition" mount -o rw $mount

    return $rc
}

# -----------------------------------------------------------------------------
function prepare_exclusion_list() {
    local uboot_path=$1

    local xlist=($uboot_path) # Exclude u-boot from main extraction

    # Write exclusion list to file, return filename
    local templist
    templist=$(mktemp -t temp$$.XXXXXX)
    printf "%s\n" "${xlist[@]}" > $templist
    echo "$templist"
}

# -----------------------------------------------------------------------------
# Escapes given string for use as a sed keyword
function sed_escape() {
    printf "%s" "$1" | sed 's/[]\/$*.^|[]/\\&/g'
}

# -----------------------------------------------------------------------------
function update_manifest() {
    local manifest_file=$1
    local xtract_list=$2

    # Remove excluded files from manifest
    while read x;
    do
        sed_x=$(sed_escape "$x")
        sed -i "/$sed_x/d" ${manifest_file}
    done < ${xtract_list}
}

# -----------------------------------------------------------------------------
manifest_add_blackmark() {
    # Add an entry to the manifest that will never exist.  This marks the
    # bank as invalid until the manifest file is fixed (used for two-part
    # upgrades)
    echo "00000000000000000000000000000000 *./blackmark/" >> ./MANIFEST
    echo "0000000000000000000000000000000000000000000000000000000000000000 *./blackmark/" >> ./MANIFEST_SHA256
}

# -----------------------------------------------------------------------------
manifest_remove_blackmark() {
    sed -i "/\*\.\/blackmark/d" ./MANIFEST
    sed -i "/\*\.\/blackmark/d" ./MANIFEST_SHA256
}

# -----------------------------------------------------------------------------
# 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 ;;
        "--import") return 0;;
        *) break ;;
    esac
done

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

active_image="$1"
src_image="$2"
src2_image="$3"
invalidate_manifest="no"

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

case "$active_image" in
    a|A) active_image="A"
         dst_app_mount=/mnt/imageA
         ;;
    b|B) active_image="B"
         dst_app_mount=/mnt/imageB
         ;;
    *)   active_image=""
         ;;
esac

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

if [ ! -f "$src_image" ] ; then
    die $EXIT_NO_SOURCE "ERROR: $src_image missing"
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

case "$src2_image" in
    *.gz)
         extract2_cmd="tar -xzpf $src2_image"
         ;;
    *.xz)
         extract2_cmd="extract_lzma_tarball $src2_image"
         ;;
    *) ;;
esac

if [ "$debug" -eq "1" ] ; then
    echo "active_image  = $active_image"
    echo "dst_app_mount = $dst_app_mount"
    echo "uboot_path    = $uboot_path"
    echo "extract_cmd   = $extract_cmd"
fi

case "$src_image" in
    *.3916.tar.gz | \
    *.3930.tar.gz | \
    *.3931.tar.gz | \
    *.3932.tar.gz | \
    *.3942.tar.gz | \
    *.5142.tar.gz | \
    *.5160.tar.gz )
        install_board_specific_files "$dst_app_mount" "$extract_cmd"
        exit $EXIT_SUCCESS
        ;;
    *)
        case "$(get_board_name)" in
            "3916" | \
            "3930" | \
            "3931" | \
            "3932" | \
            "3942" | \
            "5142" | \
            "5160" )
                # For these platforms, we are not done until a second file
                # is installed, so make sure that the manifest reflects a
                # missing part.
                invalidate_manifest="yes" ;;
            "3922" | \
            "3924" | \
            "3926" | \
            "3928" )
                # Re-use the "extract2_cmd" to populate the U-Boot and
                # initramfs staging directory if needed.
                extract2_cmd="ramfs_fixup ${dst_app_mount}" ;;
            *)
                ;;
        esac
        ;;
esac

if false ; then
    # Save the current vm dirty_bytes/dirty_ratio values.
    read vm_dirty_bytes_val < "${vm_dirty_bytes_file}" || vm_dirty_bytes_val=0
    read vm_dirty_ratio_val < "${vm_dirty_ratio_file}" || vm_dirty_ratio_val=0
    # Set vm_dirty_bytes to the lowest possible value.
    vm_dirty_shrink
fi

# Main protect logic
#
grep " $dst_app_mount " /proc/mounts \
    || exit_on_error "mount app"       mount $dst_app_mount
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}/*

# extract the rest of the files
exit_on_error "extract new files"      $extract_cmd
exit_on_error "remove fpga files"      rm -rf /mnt/sysfs/fpga/*
# sometimes need 2nd source file (like for 5160)
if [ -n "${extract2_cmd}" ]; then
   exit_on_error "extract additional files"	$extract2_cmd
fi

exit_on_error "update bootloader"      update_bootloader $dst_app_mount

# flush the cache to ensure the sha256sum 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 "sha256sum check"        sha256sum_check

if [ "$invalidate_manifest" == "yes" ] ; then
    exit_on_error "remount app rw"         mount $dst_app_mount -o remount,rw
    exit_on_error "mark partial manifest"  manifest_add_blackmark
    exit_on_error "remount app ro"         mount $dst_app_mount -o remount,ro
fi

exit_on_error "syncing app"            sync

# Restore the original vm dirty values.
vm_dirty_restore

exit $EXIT_SUCCESS
