#!/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
exp_bank_size=144 # size (in MiB) for expanded app bank

# 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

XGRADE_IMAGE_PATH=/tmp/expansion
LOGFILE=/mnt/sysfs/fault/expansion.log

INIT_PID=1                # PID of init process


# -----------------------------------------------------------------------------
print_usage_and_exit() {
    echo "Usage: $this_script [option] <A|B> <tarball> [optional tarball]"
    echo "   or: $this_script --write-only <app_mnt> <kern_mount> <uboot> <uboot-path>"
    echo "               <xlist> <dirty_bytes> <dirty_ratio> <extract_cmd> [<extract2_cmd>]"
    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
}

# -----------------------------------------------------------------------------
get_board_name_brego()
{
    board3940=/sys/bus/platform/devices/board_id.1
    board3960=/sys/bus/platform/devices/board_id.2

    if [ -e "$board3940" ]; then
        echo '3940'
    elif [ -e "$board3960" ]; then
        echo '3960'
    else
        echo 'unknown'
    fi
}

# -----------------------------------------------------------------------------
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_dernhelm()
{
    local model_file="/proc/device-tree/model"

    if [ -f "$model_file" ] ; then
        cat $model_file
    else
        echo 'unknown'
    fi
}

# -----------------------------------------------------------------------------
get_board_name()
{
    case "$(cat /family)" in
        "brego")    get_board_name_brego    ;;
        "caliondo") get_board_name_caliondo ;;
        "dernhelm") get_board_name_dernhelm ;;
        *) 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 "$?"
}

# -----------------------------------------------------------------------------
log_msg()
{
    local _level="$1"
    local _msg="$2"

    echo "${_level} ${_msg}" 
}

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

        # clean up any expansion files
        command_wrapper "clean up expansion files" rm -rf $XGRADE_IMAGE_PATH/* 
        command_wrapper "clean up inittab" sed -i '/software_expansion/d' /etc/inittab 

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

        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" -o \
         "${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
    cat $image | xz -d -c | tar -xpf - $*
}

# -----------------------------------------------------------------------------
list_lzma_tarball() {
    local image=$1; shift
    cat $image | xz -d -c | 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
}

# -----------------------------------------------------------------------------
sha256sum_check() {
    # run sha256sum check on the files in the manifest
    sha256sum -s -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=$(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)

    [ $? == 0 ] || die $EXIT_COMMAND_FAILED "Failed to find ubi volume $volume"

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

}

# -----------------------------------------------------------------------------
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 extract_uboot() {
    $extract_cmd -O $uboot_path | dd bs=128k of=$dst_uboot
}

# -----------------------------------------------------------------------------
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=$(mktemp -t temp$$.XXXXXX)
    printf "%s\n" "${xlist[@]}" > $templist
    echo "$templist"
}

# -----------------------------------------------------------------------------
# Escapes given string for use as a sed keyword
function sed_escape() {
    printf "$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}
}

# -----------------------------------------------------------------------------
mbytes_to_leb_count()
{
    local _mbytes=$1
    local _leb_size=$(ubinfo -a | awk '/Logical eraseblock size:/ { print $4}')
    local _leb_count=0

    if [[ $_mbytes =~ [0-9]+$ ]] ; then
        _leb_count=$((_mbytes * 1024 * 1024 / _leb_size))
        _leb_count=$((_leb_count + 1))
    fi

    echo $_leb_count
} 

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

# -----------------------------------------------------------------------------
create_board_specific_script() {

dst_app_mount=$1
extract_cmd=$2
vm_dirty_bytes_val=$3
vm_dirty_ratio_val=$4

    (
    cat <<EOF
# during this operation, if there is a screwup boot back to active bank
/ciena/scripts/bankset \$active_bank
$XGRADE_IMAGE_PATH/install_board_specific_files $dst_app_mount "$extract_cmd" $vm_dirty_bytes_val $vm_dirty_ratio_val
if [ "\$?" -ne "$EXIT_SUCCESS" ] ; then
    exit "\$?"
fi
/ciena/scripts/bankset \$target_bank

EOF
    ) >> $XGRADE_IMAGE_PATH/software_expansion


    (
    cat <<"EOF"
#!/bin/bash
#

#globals
debug=0
dry_run=0

LOGFILE=/mnt/sysfs/fault/expansion.log
# vm dirty byte counts
vm_dirty_bytes_file="/proc/sys/vm/dirty_bytes"
vm_dirty_ratio_file="/proc/sys/vm/dirty_ratio"


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

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

    command_wrapper "$@"

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

        # 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

        vm_dirty_restore

        exit $EXIT_COMMAND_FAILED
    fi
}

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

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

# -----------------------------------------------------------------------------
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" >> $LOGFILE
        echo "$command_line" >> $LOGFILE
        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 "$?"
}

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

# -----------------------------------------------------------------------------
sha256sum_check() {
    # run sha256sum check on the files in the manifest
    sha256sum -s -c ./MANIFEST_SHA256
}

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

# -----------------------------------------------------------------------------
# $1 -- directory $2- install command

dir=$1
cmd=$2
vm_dirty_bytes_val=$3
vm_dirty_ratio_val=$4
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

EOF
    ) > $XGRADE_IMAGE_PATH/install_board_specific_files

    chmod 755 $XGRADE_IMAGE_PATH/install_board_specific_files
}

# -----------------------------------------------------------------------------
create_shutdown_script() {

    this_script=$1
    targ_size=$2
    targ_bank=$3
    dst_app_mount=$4
    dst_kern_mount=$5
    dst_uboot=$6
    uboot_path=$7
    tempxlist=$8
    vm_dirty_bytes_val=$9
    vm_dirty_ratio_val=${10}
    extract_cmd=${11}
    extract2_cmd=${12}

storexlist="$XGRADE_IMAGE_PATH/tempxlist"

# need to copy tempxlist for safe keeping
cp $tempxlist $storexlist
tempxlist=$storexlist

    (
    cat <<EOF
#!/bin/bash

LOGFILE=/mnt/sysfs/fault/expansion.log
# during this operation, if there is a screwup boot back to active bank
/ciena/scripts/software_bank > /dev/null
running_bank="\$?"
if [ "\$running_bank" = "1" ] ; then
    active_bank="a"
    target_bank="b"
else
    active_bank="b"
    target_bank="a"
fi
/ciena/scripts/bankset \$active_bank
$XGRADE_IMAGE_PATH/expand_image --force --size $targ_size $targ_bank
if [ "\$?" -ne "$EXIT_SUCCESS" ] ; then
    exit "\$?"
fi
$XGRADE_IMAGE_PATH/software_install --write-only $dst_app_mount $dst_kern_mount $dst_uboot $uboot_path $tempxlist $vm_dirty_bytes_val $vm_dirty_ratio_val "$extract_cmd" "$extract2_cmd"
if [ "\$?" -ne "$EXIT_SUCCESS" ] ; then
    exit "\$?"
fi
/ciena/scripts/bankset \$target_bank

EOF
    ) > $XGRADE_IMAGE_PATH/software_expansion

    cp $this_script $XGRADE_IMAGE_PATH/software_install

    (
    cat <<"EOF"
#!/bin/bash
#
# This script is designed to expand an image bank by taking space away
# from the log partition.  It needs to be run prior to mounting the
# partitions that are to be resized.  The easiest way to do this is to
# boot up a system using the nosaos flag.  This script may also work on
# an NFS mounted system, but that has not been tested.
#
# This is intended to be a designer tool to allow installation of images
# that are larger than the default bank size (for prototyping purposes).
#

# globals
debug=0       # enabled by --debug or --dry-run
dry_run=0     # enabled by --dry-run
force=0       # enabled by --force
targ_size=160 # override with --size

source /ciena/scripts/utils.sh

LOGFILE=/mnt/sysfs/fault/expansion.log

EXIT_SUCCESS=0
EXIT_ERROR=1
EXIT_INVALID_ARGS=2
EXIT_NOT_ROOT=3
EXIT_COMMAND_FAILED=4

UBIFS_MIN_LEB_COUNT=21      # Minimum LEB count for UBIFS
MONITOR_DISABLE_RESPAWN=4
MONITOR_SIGTERM=2

# -----------------------------------------------------------------------------
print_usage_and_exit() {
    echo
    echo "Usage: $this_script [options] <A|B>"
    echo
    echo -e "\tThis script will expand the target image partition."
    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--size <x> provides an alternate target size in MB (default=$targ_size)"
    echo -e "\t--force    if processes still mount on /mnt/log, will kill them"
    echo -e "\t--help     shows this message"
    echo
    exit $EXIT_INVALID_ARGS
}

# -----------------------------------------------------------------------------
log_msg()
{
    local _level="$1"
    local _msg="$2"

    echo "${_level} ${_msg}" >> $LOGFILE
}

# -----------------------------------------------------------------------------
error_exit()
{
    local _msg="$1"
    local _exit_code="$2"

    if [ -z "$_exit_code" ] ; then
        _exit_code="$EXIT_ERROR"
    fi

    log_msg ERROR "$_msg"

    exit $_exit_code
}

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

    log_msg INFO "$comment"

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

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

    $command_line
}

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

    command_wrapper "$@"

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

        # recovery code would go here, if there was any

        exit $EXIT_COMMAND_FAILED
    fi
}

# -----------------------------------------------------------------------------
ubifs_is_mounted()
{
    mount | grep -q ubi0:$1
}

# -----------------------------------------------------------------------------
ubifs_mount_point()
{
    grep ubi0:$1 /etc/fstab | awk '{ print $2 }'
}

# -----------------------------------------------------------------------------
resize_ubifs_partition()
{
    local _name="$1"      # ubifs volume name
    local _leb_targ="$2"  # new leb count
    local _ro_mount="$3"  # mount point is read-only

    local _mount_point=$(ubifs_mount_point $_name)

   # check that partition is not in use
    if ubifs_is_mounted $_name ; then
        umount $_mount_point 2> /dev/NULL
        umount_ret="$?"
        if [ "$umount_ret" -ne "0" ] ; then
            if [ "$force" -eq "1" ] ; then
                # need to set guardian to play-nice in order to avoid watchdog reset on 5150
                command_wrapper "guardian play-nice" /ciena/bin/guardian play-nice
                # tell monitor to send SIGTERM and disable respawn
                echo -n $MONITOR_DISABLE_RESPAWN > /tmp/fifo/mon/saos
                echo -n $MONITOR_SIGTERM > /tmp/fifo/mon/saos
                until [ "$umount_ret" -eq "0" ]; do
                    curr_pid=$(lsof $_mount_point | awk 'NR == 2 {print $2}')
                    if [ ! -z "$curr_pid" ] ; then
                        kill -s SIGKILL $curr_pid
                    fi
                    umount $_mount_point 2> /dev/NULL
                    umount_ret="$?"
                done
            else
                log_msg INFO "current users via lsof:"
                lsof $_mount_point
                error_exit "could not unmount $_mount_point"
            fi
        fi
    fi

    # create a tmpfs to hold volume contents and copy the guts over

    exit_on_error "mount $_mount_point" mount $_mount_point

    tmp_size=$(du -s $_mount_point | cut -f1)   # grab the size
    tmp_size=$((tmp_size+100))                  # add some extra

    exit_on_error "mkdir /tmp/space" mkdir -p /tmp/space

    exit_on_error "mount /tmp/space ${tmp_size}K" \
        mount -t tmpfs tmpfs /tmp/space -o size=${tmp_size}K

    exit_on_error "copy $_mount_point to /tmp/space" \
        cp -a ${_mount_point}/. /tmp/space

    # resize and format ubifs volume
    
    exit_on_error "unmount ${_mount_point}" umount ${_mount_point}

    exit_on_error "resize ${_name} volume" \
        ubirsvol /dev/ubi0 -N ${_name} --lebs=${_leb_targ}

    local vol_id=$(ubinfo /dev/ubi0 --name ${_name} \
        | awk '/^Volume ID:/ {print $3}')

    exit_on_error "empty ubi volume" \
        ubiupdatevol -t /dev/ubi0_${vol_id} 
    exit_on_error "format ubifs for ${_name} volume" \
        mkfs.ubifs --log-lebs=3 /dev/ubi0_${vol_id}

    exit_on_error "mount ${_mount_point}" mount ${_mount_point} -o rw

    exit_on_error "copy to ${_mount_point}" cp -a /tmp/space/. ${_mount_point}

    if [ -n ${_ro_mount} ] ; then
        exit_on_error "mount ${_mount_point} ro" \
            mount ${_mount_point} -o remount,ro
    fi

    exit_on_error "unmount /tmp/space" umount /tmp/space

    exit_on_error "remove /tmp/space" rm -r /tmp/space
    
}

# -----------------------------------------------------------------------------
shrink_log_partition()
{
    local _leb_steal="$1"  # how many blocks to take from the log partition
    local _leb_thresh=600  # how many blocks to reserve for the log partition
    local _leb_curr=$(ubinfo /dev/ubi0 -N logs | awk '/^Size:/ {print $2}')
    local _leb_targ=$((_leb_curr - _leb_steal))

    log_msg INFO "reducing log partition ($_leb_curr) by $_leb_steal blocks"

    if [ "$_leb_targ" -lt "$_leb_thresh" ] ; then
        # if there is a /mnt/xftp partition, try to use it
        if ubifs_is_mounted xftp ; then
            _leb_thresh=$UBIFS_MIN_LEB_COUNT  # reduce xftp partition to this
            _leb_curr=$(ubinfo /dev/ubi0 -N xftp | awk '/^Size:/ {print $2}')
            _leb_targ=$((_leb_curr - _leb_steal))

            if [ "$_leb_targ" -lt "$_leb_thresh" ] ; then
                error_exit "not enough free blocks in xftp partition ($_leb_targ < $_leb_thresh)"
            else
                exit_on_error "resizing xftp partition" resize_ubifs_partition "xftp" "${_leb_thresh}"
                return
            fi
        else
            error_exit "not enough free blocks in log partition ($_leb_targ < $_leb_thresh)"
        fi
    fi

    # If the partition isn\'t mounted, mount it so that we can clear core files.
    if ! ubifs_is_mounted logs ; then
        exit_on_error "mount /mnt/log" mount /mnt/log
    fi

    exit_on_error "clear /mnt/log/corefiles" \
        rm -rf /mnt/log/corefiles/*

    exit_on_error  "resizing log partition" resize_ubifs_partition "logs" "${_leb_targ}"
}

# -----------------------------------------------------------------------------
mbytes_to_leb_count()
{
    local _mbytes=$1
    local _leb_size=$(ubinfo -a | awk '/Logical eraseblock size:/ { print $4}')
    local _leb_count=0

    if [[ $_mbytes =~ [0-9]+$ ]] ; then
        _leb_count=$((_mbytes * 1024 * 1024 / _leb_size))
        _leb_count=$((_leb_count + 1))
    fi

    echo $_leb_count
} 

# -----------------------------------------------------------------------------
get_board_name_pretty()
{
    board_name=$(get_board_name)

    # Handle special cases where the same board id is used for multiple
    # board variants.  Regular cases get the unmodified result from
    # the get_board_name result.
    #
    case "$board_name" in

        "3932")
            if ! pwe_supported ; then
                echo "3930ET"
                return
            fi
            ;;

        "3916")
            if cr3916 ; then
                echo "3916CR"
                return
            fi
            ;;

    esac

    echo $board_name
}

# -----------------------------------------------------------------------------
# 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 ;;
        "--size") shift ; targ_size=$1 ; shift ;;
        "--force") force=1 ; shift ;;
        *) break ;;
    esac
done

targ_leb_count=$(mbytes_to_leb_count $targ_size)

if [ "$targ_leb_count" -eq "0" ] ; then
    error_exit "invalid target size  \"$targ_size\"" $EXIT_INVALID_ARGS
fi

# get positional arguments
target_bank=$1

if [ "$#" -eq "0" ] ; then
    log_msg ERROR "missing target image partition"
    print_usage_and_exit
fi

case ${target_bank} in
    'A') standby_bank='B' ;;
    'B') standby_bank='A' ;;
    *) error_exit "invalid target bank \"$target_bank\"" $EXIT_INVALID_ARGS
esac

# check that we can proceed:
# - product family must be caliondo
# - must not resize the running rootfs
# - must be running as root

family=$(cat /family)
model=$(get_board_name_pretty)

if [ "$family" != "caliondo" ] ; then
    error_exit "can not run on $family $model (only supported on caliondo)"
fi

rootfs_bank=$(kernel_arg root)
rootfs_bank=${rootfs_bank#ubi0:app}

if [ "$rootfs_bank" == "$target_bank" ] ; then
    error_exit "can not resize running rootfs" $EXIT_INVALID_ARGS
fi

if [ "$(id -u)" -ne "0" ] ; then
    error_exit "must be running as root" $EXIT_NOT_ROOT
fi

# check size of target bank
curr_leb_count=$(ubinfo --devn=0 --name=app${target_bank} \
    | awk '/Size/ {print $2}')
standby_leb_count=$(ubinfo --devn=0 --name=app${standby_bank} \
    | awk '/Size/ {print $2}')

if [ "$curr_leb_count" -ge "$targ_leb_count" ] && [ "$standby_leb_count" -ge "$targ_leb_count" ] ; then
    log_msg INFO "active_size=$curr_leb_count, standby_size=$standby_leb_count targ=$targ_leb_count, no increase required"
    exit $EXIT_SUCCESS
fi

# resize logic starts here

free_leb_count=$(ubinfo -a \
    | awk '/Amount of available logical eraseblocks:/ { print $6}')

extr_leb_count=$((targ_leb_count * 2 - curr_leb_count - standby_leb_count - free_leb_count))

# If there are already enough unused blocks, then just use them.  Otherwise
# steal blocks from the log partition to be used in the app partition.

if [ "$extr_leb_count" -gt "0" ] ; then
    exit_on_error "shrinking log partition" shrink_log_partition $extr_leb_count
else
    log_msg INFO "only free blocks will be used for increase"
fi

# Only resize standby image for now - active gets done after reboot
exit_on_error "resizing image partition" resize_ubifs_partition "app${target_bank}" "${targ_leb_count}" "yes"

log_msg INFO "successfully expanded app${target_bank}"

exit $EXIT_SUCCESS
EOF
    ) > $XGRADE_IMAGE_PATH/expand_image

    chmod 755 $XGRADE_IMAGE_PATH/software_expansion
    chmod 755 $XGRADE_IMAGE_PATH/software_install
    chmod 755 $XGRADE_IMAGE_PATH/expand_image

}

# -----------------------------------------------------------------------------
software_write() {

dst_app_mount=$1
dst_kern_mount=$2
dst_uboot=$3
uboot_path=$4
tempxlist=$5
vm_dirty_bytes_val=$6
vm_dirty_ratio_val=$7
extract_cmd=$8
extract2_cmd=$9

# 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_format $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}/*

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}/*
    #
    # Aim ${dst_app_mount}/boot straight at ${dst_kern_mount}. And
    # remember to update the soft link between the alternate banks in
    # the software_protect script.
    #
    # The tar extract_cmd will redirect all files under ./boot to the
    # /mnt/kernel[AB] partition - which is the desired location on
    # brego and NOR-flash-based caliondo boards.
    #
    # The soft link also leaves the manifest file paths intact, and it
    # keeps happy the FPGA installers that expect images under ./boot.
    #
    exit_on_error "add boot soft link" ln -s ${dst_kern_mount} ${dst_app_mount}/boot
fi

# extract the rest of the files
exit_on_error "extract new files"      $extract_cmd -X $tempxlist
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

if copy_kernel; then
    exit_on_error "remount kernel ro"  mount ${dst_kern_mount} -o remount,ro
fi

# Exclude the U-Boot image from the manifest files
update_manifest ./MANIFEST_SHA256 ${tempxlist}
update_manifest ./MANIFEST ${tempxlist}

# Finally remove the exclusion list file
rm -f $tempxlist

# 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 "extract u-boot"         extract_uboot
exit_on_error "syncing app"            sync

# Restore the original vm dirty values.
vm_dirty_restore
}

# -----------------------------------------------------------------------------
# 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 ;;
        "--write-only") software_write $2 $3 $4 $5 $6 $7 $8 "$9" "${10}"; exit $EXIT_SUCCESS; shift ;;
        *) 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"
         standby_image="B"
         dst_app_mount=/mnt/imageA
         dst_kern_mount=/mnt/kernelA
         dst_uboot=/dev/$(find_mtd_device u-boot-a "Uboot A")
         if [ "$?" -ne "0" ] ; then
             die $EXIT_COMMAND_FAILED "Unable to locate mtd device for u-boot a"
         fi
         ;;
    b|B) active_image="B"
         standby_image="A"
         dst_app_mount=/mnt/imageB
         dst_kern_mount=/mnt/kernelB
         dst_uboot=/dev/$(find_mtd_device u-boot-b "Uboot B")
         if [ $? -ne 0 ]; then
             die $EXIT_COMMAND_FAILED "Unable to locate mtd device for u-boot b"
         fi
         ;;
    *)   active_image=""
         ;;
esac

# Boards that boot from SPI flash use a special U-Boot image.
if spiboot_supported; then
    uboot_path=./ciena/images/u-boot-spi.img
else
    uboot_path=./ciena/images/u-boot.img
fi

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 "dst_kern_mount= $dst_kern_mount"
    echo "uboot_path    = $uboot_path"
    echo "dst_uboot     = $dst_uboot"
    echo "extract_cmd   = $extract_cmd"
fi

family=$(cat /family)
targ_leb_count=0
curr_leb_count=0
standby_leb_count=0

if [ "$family" == "caliondo" ] ; then
    targ_leb_count=$(mbytes_to_leb_count $exp_bank_size)
    curr_leb_count=$(ubinfo --devn=0 --name=app${active_image} \
         | awk '/Size/ {print $2}')
    standby_leb_count=$(ubinfo --devn=0 --name=app${standby_image} \
         | awk '/Size/ {print $2}')
fi 

# 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

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 )
        if [ "$family" == "caliondo" ] ; then
             if [ "$targ_leb_count" -gt "$curr_leb_count" ] || [ "$targ_leb_count" -gt "$standby_leb_count" ] ; then
                 # copy image files to different directory so they don't get deleted
                 src_name=$(basename "$src_image")
                 exit_on_error "make directory to hold image" mkdir -p $XGRADE_IMAGE_PATH
                 src_backup_image=$XGRADE_IMAGE_PATH/$src_name
                 exit_on_error "copy board-specific image file" cp $src_image $src_backup_image
                 extract_backup_cmd=$(echo "$extract_cmd" | sed  "s+$src_image+$src_backup_image+" )
       
                 create_board_specific_script "$dst_app_mount" "$extract_backup_cmd" $vm_dirty_bytes_val $vm_dirty_ratio_val
            else
                 install_board_specific_files "$dst_app_mount" "$extract_cmd"
            fi
        else
            install_board_specific_files "$dst_app_mount" "$extract_cmd"
        fi

        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" ;;
            *)
                ;;
        esac
        ;;
esac

# If the force option is specified, then unmount the destination partitions
# and erase them.
#
if [ "$force" -eq "1" ] ; then
    if copy_kernel; then
        erase_partition "kernel" ${dst_kern_mount} \
            || die $EXIT_COMMAND_FAILED "failed to erase kernel partition"
    fi

    erase_partition "app"  ${dst_app_mount}  \
            || die $EXIT_COMMAND_FAILED "failed to erase app partition"
fi

tempxlist=$(prepare_exclusion_list $uboot_path)

# if running on Caliondo and image bank needs to be expanded
family=$(cat /family)

if [ "$family" == "caliondo" ] ; then
    if [ "$targ_leb_count" -gt "$curr_leb_count" ] || [ "$targ_leb_count" -gt "$standby_leb_count" ] ; then
        log_msg INFO "attempting to resize app${active_image} and app${standby_image} to ${targ_leb_count} blocks (${exp_bank_size} MB)"
        # add line to  inittab so that our file will be executed on shutdown
        echo "::shutdown:$XGRADE_IMAGE_PATH/software_expansion" >> /etc/inittab
        # send signal to BusyBox init to pick up inittab changes
        kill -s SIGHUP $INIT_PID

        exit_on_error "mkdir $XGRADE_IMAGE_PATH" mkdir -p $XGRADE_IMAGE_PATH
        exit_on_error "mount $XGRADE_IMAGE_PATH" \
            mount -t tmpfs tmpfs $XGRADE_IMAGE_PATH
        # copy image files to different directory so they don't get deleted
        src_name=$(basename "$src_image")
        src_backup_image=$XGRADE_IMAGE_PATH/$src_name
        exit_on_error "copy initial image file" cp $src_image $src_backup_image
        extract_backup_cmd=$(echo "$extract_cmd" | sed  "s+$src_image+$src_backup_image+" )
       
        if [ ! -z "$src2_image" ] ; then
            src2_name=$(basename "$src2_image")

            src2_backup_image=$XGRADE_IMAGE_PATH/$src2_name
            exit_on_error "copy board-specific image file" cp $src2_image $src2_backup_image
            extract2_backup_cmd=$(echo "$extract2_cmd" | sed  "s+$src2_image+$src2_backup_image+" )
        fi
 
        # create script file to run on shutdown
        create_shutdown_script $0 $exp_bank_size $active_image $dst_app_mount $dst_kern_mount $dst_uboot $uboot_path $tempxlist $vm_dirty_bytes_val $vm_dirty_ratio_val "$extract_backup_cmd" "$extract2_backup_cmd"
        exit $EXIT_SUCCESS
    else
        log_msg INFO "active_size=$curr_leb_count, standby_size=$standby_leb_count, targ=$targ_leb_count, no increase required"
    fi
fi

software_write $dst_app_mount $dst_kern_mount $dst_uboot $uboot_path $tempxlist $vm_dirty_bytes_val $vm_dirty_ratio_val "$extract_cmd" "$extract2_cmd"

exit $EXIT_SUCCESS
