##
## Copyright 2013 Cumulus Networks Inc.
## All rights reserved
##
## Installer functions and utilities
##
## This assumes the shell is ash/dash.
##
## Assume we are in an environment that supplies functions like:
##   cl_run
##   log_failure_msg
##   log_success_msg
##   log_begin_msg
##   log_end_msg
##
## These functions support a number of operations on Cumulus image
## files, including:
##
##   verify  -- Verify image integrity and return.
##   extract -- Extract constituent files, no installation performed.
##   info    -- Dump image meta data and file TOC.
##   install -- Extract and install image files.
##   env     -- Extract and install u-boot environment variables.
##
## On error the functions herein return 1 and do not exit.  The
## callers of these functions need to check the status codes and clean
## things up themselves.

# Trap CTRL-C
cl_trap_ctrl_c()
{
    trap '' INT
}

# Trap CTRL-C, CTRL-Z and CTRL-\
cl_trap_signals()
{
    trap '' INT TSTP QUIT
}

# Enable CTRL-C
cl_enable_ctrl_c()
{
    trap - INT
}

# Enable CTRL-C, CTRL-Z and CTRL-\
cl_enable_signals()
{
    trap - INT TSTP QUIT
}

_calc_blks()
{
    sz=$1
    blk_sz=$2
    n_blks=$(( ( $sz + $blk_sz - 1 ) / $blk_sz ))
    echo -n $n_blks
}

##
## Check existence and readability of control file.
##
_check_control()
{
    if [ ! -f control ] ; then
        log_failure_msg "Missing control file."
        return 1
    fi
    if [ ! -r control ] ; then
        log_failure_msg "Unreadable control file."
        return 1
    fi
}

##
## Verify the *.sha1 files contained in control.tar.gz.  During
## testing they may not be present, so only check them if they exist.
##
_verify_control_sha1()
{
    # Verify sha1's if present
    sha_files=$(ls *.sha1 2> /dev/null)
    for f in $sha_files ; do
        base=${f%.sha1}
        log_debug "Validating sha1 for $base ..."
        $SHA1SUM -c ${base}.sha1 > /dev/null 2>&1
        result=$?
        if [ $result -ne 0 ] ; then
            log_failure_msg "Unable to verify sha1 for $base."
            return 1
        fi
    done
}

##
## Check the Installer-Version: variable from the control.  Depending
## on how we're called either print a warning or print an error.
##
_check_installer_version()
{
    on_mismatch=$1

    # verify the "Installer-Version" variable.
    installer_ver=$(grep 'Installer-Version:' control | awk '{ print $2 }')
    expected_ver="1.2"
    if [ "$installer_ver" != "$expected_ver" ] ; then
        if [ "$on_mismatch" = "error" ] ; then
            log_failure_msg "Unexpected Installer-Version $installer_ver.  Expecting \`$expected_ver' ."
            return 1
        else
            log_warning_msg "Unexpected Installer-Version $installer_ver.  Expecting \`$expected_ver' ."
            return 0
        fi
    fi
}

##
## image info operation handler
##
image_info_handler()
{
    _check_control
    _verify_control_sha1
    _check_installer_version warn

    _log_msg "Control File Contents\n"
    _log_msg "=====================\n"
    # Dump the control file, without printing the supported platforms
    # field.
    sed -e '/^Platforms:/d' control
    _log_msg "\n"
    _log_msg "Data Archive Contents\n"
    _log_msg "=====================\n"
    $TAR tvf $data_tar || {
        log_failure_msg "Errors dumping archive file list."
        return 1
    }
}

_data_sha1_verify()
{
    if [ $# -ne 0 ] ; then
        sha1_files=
        for f in "$@" ; do
            sha1_files="$sha1_files $f.sha1"
        done
    else
        # get the list of sha1 files in data_tar
        sha1_files="$($TAR tf $data_tar | grep sha1)"
    fi

    # compare sha1 file to computed sha1
    err=0
    for f in $sha1_files ; do
        base=${f%.sha1}
        log_begin_msg "Validating sha1 for $base"
        # extract .sha1 file
        $TAR xf $data_tar $f || {
            log_failure_msg "Unable to extract sha1 file $f."
            return 1
        }
        # compare sha1 file to what is in archive
        $TAR -xOf $data_tar $base | $SHA1SUM -c $f > /dev/null 2>&1
        result=$?
        log_end_msg
        if [ $result -ne 0 ] ; then
            log_failure_msg "Unable to verify sha1 for $base."
            err=$(( $err + 1 ))
        fi

        # clean up
        rm -f $f
    done

    if [ $err -ne 0 ] ; then
        return 1
    fi
}

##
## image verify operation handler
##
image_verify_handler()
{
    _check_control || return 1
    _verify_control_sha1 || return 1
    _check_installer_version error || return 1
    _data_sha1_verify || return 1
}

##
## image extract operation handler
##
image_extract_handler()
{
    [ -z "$1" ] && {
        log_failure_msg "The 'extract' verb requires an output directory argument"
        return 1
    }
    # check for absolute path
    [ "$(echo $1 | head -c 1)" = "/" ] || {
        log_failure_msg "The output directory must be an absolute path: $1"
        return 1
    }
    output_dir=$(realpath $1) || {
        log_failure_msg "Problems with output directory: $1"
        return 1
    }
    _check_control
    _verify_control_sha1
    _check_installer_version warn

    [ -d "$output_dir" ] || {
        log_failure_msg "Output directory does not exist: $output_dir"
        return 1
    }
    log_debug "Extracting files from archive into $output_dir ...\n"
    $TAR -C $output_dir -xvf $data_tar || {
        log_failure_msg "Unable to extract archive files from $data_tar."
        return 1
    }
}

# Return the default kernel file name(s) to install
get_kernel_file()
{
    echo -n "$default_kernel_file"
}

# Return the default sysroot file name to install
get_sysroot_file()
{
    echo -n "$default_sysroot_file"
}

##
## image install operation handler
##
image_install_handler()
{

    log_pre=
    if [ "$dry_run" = "y" ] ; then
        log_pre="Dry run: "
    fi

    _check_control || return 1
    _verify_control_sha1 || return 1
    _check_installer_version error || return 1

    # verify the current architecture is equal to the "Architecture" variable.
    installer_arch=$(grep 'Architecture:' control | awk '{ print $2 }')
    if [ "$installer_arch" != "$arch" ] ; then
        log_failure_msg "Unexpected Architecture $installer_arch.  Expecting \`$arch'."
        return 1
    fi
    log_debug "Installing image for architecture: $installer_arch ."

    log_debug "Installing image for platform: $platform ."

    # extract the kernel image and sha1
    kimage=$(get_kernel_file) || return 1
    sysroot=$(get_sysroot_file) || return 1

    # verify integrity of uImage and sysroot
    _data_sha1_verify $kimage $sysroot || return 1

    $TAR -xf $data_tar $kimage || {
        log_failure_msg "Unabled to extract kernel image $kimage from $data_tar."
        return 1
    }
    $TAR -xf $data_tar ${kimage}.sha1 || {
        log_failure_msg "Unabled to extract kernel image sha1 ${kimage}.sha1 from $data_tar."
        return 1
    }
    log_debug "Installing kernel image: $kimage ."

    # extract the sysroot sha1
    $TAR -xf $data_tar ${sysroot}.sha1 || {
        log_failure_msg "Unabled to extract ${sysroot}.sha1 from $data_tar."
        return 1
    }

    # extract the size of the sysroot squashfs file
    local sysroot_sz
    sysroot_sz=$($TAR tvf $data_tar $sysroot | awk '{ print $3 }')
    [ $? -eq 0 ] || {
        log_failure_msg "Unabled to determine size of $sysroot."
        return 1
    }
    log_debug "sysroot size: $sysroot_sz"

    sha1_sysroot=$(cat ${sysroot}.sha1 | awk '{ print $1 }')
    os_release=$(grep OS-Release: control | awk '{print $2}')

    max_attempts=3

    for slot in "$@" ; do
        [ $slot -ne 1 ] && [ $slot -ne 2 ] && {
            log_failure_msg "Bad image slot.  Expecting 1 or 2."
            return 1
        }
        _log_msg "${log_pre}Installing OS-Release $os_release into image slot $slot ...\n"

        # Two phase commit on the update.  When installing, first set a
        # state variable to indicate the install is in progress.  After
        # successful install set the state variable to the OK state.
        cl_slot_state_set $slot 1

        local srs
        [ "$#" -eq 2 ] && {
            # installing into 2 slots, must be from ONIE
            [ $slot -eq 1 ] && srs="50%FREE"
            [ $slot -eq 2 ] && srs="100%FREE"
        }

        # Allow for retries with bad sysroot programming
        attempt=0
        while [ $attempt -lt $max_attempts ] ; do
            install_sysroot $slot $srs && {
                log_begin_msg "${log_pre}Verifying sysroot copy"
                verify_sysroot $slot $sha1_sysroot $sysroot_sz
                if [ $? -eq 0 ]; then
                    _log_msg " OK.\n"
                    break
                else
                    _log_msg " Failed.\n"
                    log_failure_msg "${log_pre}Unable to verify image copy to sysroot slot ${slot}."
                fi
            }
            _log_msg "${log_pre}Retrying copy...\n"
            attempt=$(( $attempt + 1 ))
        done

        if [ $attempt -eq $max_attempts ] ; then
            log_failure_msg "${log_pre}Unable to program sysroot slot ${slot} after $attempt attempts."
            _log_msg "${log_pre}Giving up...\n"
            # set the slot state to error
            cl_slot_state_set $slot 2
            return 1
        fi

        kernel_install_handler $kimage $slot
        if [ $? -ne 0 ] ; then
            # set the slot state to error
            _log_msg " Failed.\n"
            log_failure_msg "${log_pre}Unable to program kernel copy to sysroot slot ${slot}."
            cl_slot_state_set $slot 2
            return 1
        fi

        # Mark slot state as OK
        cl_slot_state_set $slot 0 || return 1
        cl_slot_version_set $slot $os_release || return 1

        arch_post_slot_install $slot || return 1

        if [ "$#" -eq 2 ] && [ $slot -eq 1 ] ; then
                # perform arch dependent final processing
                # calling arch_post_provision before installing slot2
                # That way, if any problems (power outage) happens
                # during slot 2 install, the grub "cl.active" variable
                # will still be set correctly
                arch_post_provision || {
                    log_failure_msg "${log_pre}arch_post_provision failed."
                    return 1
                }
        fi

   done

}

#
# Fresh install on a system.  Runs in ONIE context.
# Perform the following:
#
#    - partition/format mass storage device
#    - install image into both image slots
#    - update uboot env variables
#    - reboot
#
image_provision_handler()
{
    log_pre=
    if [ "$dry_run" = "y" ] ; then
        log_pre="Dry run: "
    fi

    if [ $root_type = 'blk' ]; then
        # verify block device partitions
        for p in $ro_part1 $ro_part2 $rw_rootpart $persist_part; do
            [ -b /dev/$p ] || {
                log_failure_msg "${log_pre}Destination /dev/$p is not a valid block device."
                return 1
            }
        done
    fi

    if [ "$installer_mtd" = "yes" ] ; then
        # verify mtd device
        # look for "open" mtd device, found in onie dtb
        for d in /dev/mtd-open ; do
            [ -h $d ] || {
                log_failure_msg "${log_pre}Destination $d is not a valid MTD device."
                return 1
            }
            # flash_unlock may be unsupported on some boards, so ignore errors
            # here.
            if [ "$dry_run" = "y" ] ; then
                _log_msg "${log_pre}flash_unlock $d\n"
            else
                flash_unlock $d > /dev/null 2>&1
            fi
        done
    fi

    _log_msg "${log_pre}Dumping image info...\n"
    image_info_handler "$@" || {
        log_failure_msg "${log_pre}Dumping image info failed."
        return 1
    }

    ##
    ## Point of no return.  Any failures below and the system is wacked.
    ##
    ## Do as much error checking above before proceeding.
    ##

    # partition block devices/ubiformat mtd devices
    format_disk || {
        log_failure_msg "${log_pre}Format disk failed."
        return 1
    }

    image_install_handler 1 2 || {
        log_failure_msg "${log_pre}Image slot $slot installation failed."
        return 1
    }

    # install default environment variables
    image_env_handler || {
        log_failure_msg "${log_pre}Environment variable installation failed."
        return 1
    }

    return 0
}

##
## update/replace files within the archive and re-assemble it
##
image_update_handler()
{

    log_pre=
    if [ "$dry_run" = "y" ] ; then
        log_pre="Dry run: "
    fi

    org_archive=$1
    [ -r "$org_archive" ] || {
        log_failure_msg "${log_pre}Unable to determine original archive file."
        return 1
    }

    input_dir=$2
    [ -z $input_dir ] && {
        log_failure_msg "${log_pre}'update' operation requires an input directory argument."
        log_failure_msg "${log_pre}Make sure to pass an absolute directory path."
        return 1
    }
    [ -d $input_dir ] || {
        log_failure_msg "${log_pre}Input directory does not exist: $input_dir:"
        log_failure_msg "${log_pre}Make sure to pass an absolute directory path."
        return 1
    }

    output_file=$3
    if [ -z "$output_file" ] ; then
        # replace original file
        output_file=$org_archive
    fi

    # extract image to a temp directory
    log_begin_msg "${log_pre}Preparing image updater"
    update_tmp_dir=$(mktemp -d)
    mount -t tmpfs tmpfs-updater $update_tmp_dir
    extract="$update_tmp_dir" sh $org_archive > /dev/null 2>&1 || {
        log_end_msg
        log_failure_msg "${log_pre}Unable to extract $org_archive to $update_tmp_dir."
        umount $update_tmp_dir
        return 1
    }
    [ -d "$update_tmp_dir/installer" ] || {
        log_end_msg
        log_failure_msg "${log_pre}Directory does not exist: $update_tmp_dir/installer."
        umount $update_tmp_dir
        return 1
    }
    log_end_msg

    log_begin_msg "${log_pre}Adding files from $input_dir"
    $TAR -C $input_dir -cf - . | $TAR -C $update_tmp_dir/installer -xf - || {
        log_end_msg
        log_failure_msg "${log_pre}Problems copying files from $input_dir => $update_tmp_dir/installer."
        umount $update_tmp_dir
        return 1
    }
    log_end_msg

    # Allow for retries with bad re-packaging
    max_attempts=3
    attempt=0
    while [ $attempt -lt $max_attempts ] ; do
        log_begin_msg "${log_pre}Re-packing image"
        sharch="$update_tmp_dir/sharch.tar"
        # No need for compression here
        $TAR -C $update_tmp_dir -cf $sharch installer || {
            log_end_msg
            log_failure_msg "${log_pre}Problems creating $sharch archive"
            umount $update_tmp_dir
            return 1
        }

        [ -f "$sharch" ] || {
            log_end_msg
            log_failure_msg "${log_pre}Problems creating $sharch archive"
            umount $update_tmp_dir
            return 1
        }
        log_end_msg

        log_begin_msg "${log_pre}Calculating checksum"
        sha1=$(cat $sharch | $SHA1SUM | awk '{print $1}')
        log_end_msg

        log_begin_msg "${log_pre}Finalizing output image $output_file"
        tmp_output="$update_tmp_dir/output_file"
        sed -n -e '1,/^exit_marker$/p' $org_archive > $tmp_output || {
            log_end_msg
            log_failure_msg "${log_pre}Problems creating $tmp_output from $org_archive"
            umount $update_tmp_dir
            return 1
        }

        # Replace sha1 variable in the sharch template
        sed -i -e "s/^payload_sha1=.*$/payload_sha1=$sha1/" $tmp_output || {
            log_end_msg
            log_failure_msg "${log_pre}Problems updating checksum for $tmp_output"
            umount $update_tmp_dir
            return 1
        }
        cat $sharch >> $tmp_output
        sync; sync

        verify=y sh $tmp_output > /dev/null 2>&1 && break
        log_failure_msg "Unable to verify checksum of repackaged installer image."

        log_failure_msg "Retrying the packaging..."
        attempt=$(( $attempt + 1 ))
    done

    if [ $attempt -eq $max_attempts ] ; then
        log_failure_msg "${log_pre}Unable to repackage image after $attempt attempts."
        _log_msg "${log_pre}Giving up...\n"
        umount $update_tmp_dir
        return 1
    fi

    rm -f $output_file
    cp $tmp_output $output_file
    sync; sync
    umount $update_tmp_dir
    log_end_msg

}

# Save various information about the install to the provisioned
# system:
#   - runs in ONIE context
#   - the installer image used
#   - the log files
savedir_type=
savedir_opts=
savedir_dev=
savedir_img=no
save_info()
{
    if [ "$dry_run" = "y" ] ; then
        _log_msg "${log_pre}save_info()\n"
        return 0
    fi

    arch_set_savedir
    tmpdir=$(mktemp -d)
    mount -t $savedir_type -o $savedir_opts $savedir_dev $tmpdir || {
        echo "Error: Unable to mount $savedir_dev on $tmpdir."
        rm -rf $tmpdir
        return 1
    }
    save_dir="$tmpdir/onie"
    rm -rf $save_dir
    mkdir -p $save_dir || {
        echo "Error: Unable to create onie save directory: $save_dir"
        umount $tmpdir
        rm -rf $tmpdir
        return 1
    }
    [ "$savedir_img" = "yes" ] && cp $archive_file $save_dir/onie-installer
    arch_save_info "$save_dir"
    if [ -x /bin/onie-support ] ; then
        onie-support "$save_dir"
    else
        # gather info manually
        cp -a /var/log $save_dir
        cat /proc/cmdline > $save_dir/kernel_cmdline.txt
        export > $save_dir/runtime-export-env.txt
        set > $save_dir/runtime-set-env.txt
        ps w > $save_dir/runtime-process.txt
    fi

    umount $tmpdir
    rm -rf $tmpdir

    if [ $root_type = 'ubi' -a -d /sys/class/ubi/ubi0 ]; then
        ubidetach -d 0 /dev/ubi_ctrl
    fi
}


##
## kernel install operation handler
##
kernel_install_handler()
{

    # verify the current architecture is equal to the "Architecture" variable.
    # verify the current platform is listed in the "Platforms" variable.

    # extract the kernel image and sha1
    kimage=$1
    slot=$2

    log_debug "Installing kernel image: $kimage ."

    max_attempts=3

    [ $slot -ne 1 ] && [ $slot -ne 2 ] && {
        log_failure_msg "Bad image slot.  Expecting 1 or 2."
        return 1
    }

    # Two phase commit on the update.  When installing, first set a u-boot
    # state variable to indicate the install is in progress.  After
    # successful install set the state variable to the OK state.
    cl_slot_state_set $slot 1

    # Allow for retries with bad kernel programming
    attempt=0
    while [ $attempt -lt $max_attempts ] ; do
        install_kernel $slot $kimage && {
            # success
            break
        }
        _log_msg "Retrying kernel install...\n"
        attempt=$(( $attempt + 1 ))
    done

    if [ $attempt -eq $max_attempts ] ; then
        log_failure_msg "Unable to install kernel after $attempt attempts."
        _log_msg "Giving up...\n"
        # set the slot state to error
        cl_slot_state_set $slot 2
        return 1
    fi

    # Mark slot state as OK
    cl_slot_state_set $slot 0

    return 0
}

##
# Local Variables:
# mode: shell-script
# eval: (sh-set-shell "/bin/sh" t nil)
# End: