#!/bin/bash
#
# This script is used to clone a software image from one bank to another.
#

source /ciena/scripts/utils.sh

# 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]"
    echo
    echo -e "\tWhen called without a parameter, copies the software load from"
    echo -e "\tthe active image bank into the standby.  When called with a"
    echo -e "\tparameter, copies from the specified bank."
    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
}

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

# -----------------------------------------------------------------------------
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
        command_wrapper "mount kernel ro" mount $dst_kern_mount -o remount,ro

        exit $EXIT_COMMAND_FAILED
    fi
}

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

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

# -----------------------------------------------------------------------------
fix_boot_softlink()
{
    # the software install script sets up a ${app_mnt}/boot -> ${kern_mnt}
    # soft link: after the copy operation the soft link must be updated
    local kern_mnt=$1 app_mnt=$2

    rm -f ${app_mnt}/boot && ln -s ${kern_mnt} ${app_mnt}/boot
}

# -----------------------------------------------------------------------------
# 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 first parameter exists, then it should indicate the active image
if [ "$#" -gt "0" ] ; then
    active_image="$1"
else
    get_running_bank
    bank=$?

    case "$bank" in
        $EXIT_BANKA) active_image="A" ;;
        $EXIT_BANKB) active_image="B" ;;
        *)           active_image=""  ;;
    esac
fi

if [ "$#" -gt "1" ] ; then
    echo "ERROR: too many parameters"
    print_usage_and_exit
fi

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"
         src_app_mount=/mnt/imageA
         dst_app_mount=/mnt/imageB
         src_kern_mount=/mnt/kernelA
         dst_kern_mount=/mnt/kernelB
         src_kernel=/mnt/kernelA/*
         dst_kernel=/mnt/kernelB/
         src_uboot=/dev/mtd/u-boot-a
         dst_uboot=/dev/mtd/u-boot-b
         dst_kern_part=/dev/mtd/kernel-b
         ;;
    b|B) active_image="B"
         src_app_mount=/mnt/imageB
         dst_app_mount=/mnt/imageA
         src_kern_mount=/mnt/kernelB
         dst_kern_mount=/mnt/kernelA
         src_kernel=/mnt/kernelB/*
         dst_kernel=/mnt/kernelA/
         src_uboot=/dev/mtd/u-boot-b
         dst_uboot=/dev/mtd/u-boot-a
         dst_kern_part=/dev/mtd/kernel-a
         ;;
    *)   active_image=""
         ;;
esac

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

if [ "$debug" -eq "1" ] ; then
    echo "active_image  = $active_image"
    echo "src_app_mount = $src_app_mount"
    echo "dst_app_mount = $dst_app_mount"
    echo "src_kern_mount= $src_kern_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 "dst_kern_part = $dst_kern_part"
fi

# 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

# 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 "erase app files"        rm -rf ${dst_app_mount}/*
exit_on_error "copy app files"         cp -a ${src_app_mount}/* $dst_app_mount

if copy_kernel; then
    exit_on_error "fix boot soft link" fix_boot_softlink $dst_kern_mount $dst_app_mount
fi

exit_on_error "remount app ro"         mount $dst_app_mount -o remount,ro
exit_on_error "cd to $dst_app_mount"   cd $dst_app_mount

if copy_kernel; then
    exit_on_error "remount kernel rw"      mount $dst_kern_mount -o remount,rw
    exit_on_error "copy kernel"            cp $src_kernel $dst_kernel
    exit_on_error "remount kernel ro"      mount $dst_kern_mount -o remount,ro
fi

exit_on_error "check manifest(sha256)" sha256sum -s -c ./MANIFEST_SHA256
exit_on_error "copy u-boot"            dd bs=128k if=$src_uboot of=$dst_uboot
exit_on_error "syncing app"            sync

exit $EXIT_SUCCESS
