#!/bin/sh

#
# Copyright 2016 Cumulus Networks, Inc.
# All rights reserved.
#
# Stage an ONIE compatible Cumulus Linux installer for installation.
# Requires a reboot to take effect.
#

set -e

help_brief="[-s] [-i <installer path>] [-z <ztp path>] [-c] [-a] [-f] [-n]"
help_summary="This command stages an installer for later installation."
help_args=$(cat <<EOF

-s
        Status: Display any pending install operations.  This is the
        default if no other operations are specified.

-i
        Install: Stage an installer specified by the installer path.

-z
        Zero Touch Provisioning: Stage a Cumulus ZTP script specified
        by the ztp path.  This option can only be used in conjunction
        with the -i option.  See the example below.

-c
        Cancel: Cancel any previously staged installer.  This is used
        to undo a previous -i operation.

-a
        Activate:  Activate a previously staged installer.

	WARNING: This is a destructive operation.  Back up your data
	first.

	This option activates a previously staged installer image.  A
	reboot is required to complete the operation.

-f
        Force:  Force the operation.  Assumes "yes" to any questions the
        script might ask.

-n
        Dry run.  Show what would happen, but do not actually modify
        anything.

EOF
)
help_body=$(cat <<EOF

This command is used to stage a Cumulus Linux installer for later
installation.  The path for the installer can take the following
forms:

  local file path (absolute or relative path OK)
  http://server/path/...
  https://server/path/...
  scp://user@server/path/...
  ftp://server/path/... [anonymous only]

With the '-z' option the command will also stage a Zero Touch
Provisioning (ztp) script along with the installer.

Example 1 - Staging an installer hosted on a HTTP server:
---------------------------------------------------------

  cumulus@cumulus:~$ sudo onie-install -i http://203.0.113.10/image-installer

Note: At this point the image is only 'staged' for later install.  To
activate the installer requires using the '-a' option and entails a
reboot.

Example 2 - Staging an installer and a ztp script hosted on a HTTP server:
--------------------------------------------------------------------------

  cumulus@cumulus:~$ sudo onie-install -i http://203.0.113.10/image-installer \
                                       -z http://203.0.113.10/ztp-script

In this example the installer image is downloaded as in the previous
example.  In addition the ztp script specified by the -z <URL> is
staged along with the installer image.

If the ztp script URL does not exist the command fails and no files
are staged.

Example 3 - Activate the staged installer
-----------------------------------------

  cumulus@cumulus:~$ sudo onie-install -a
  WARNING: This will wipe out all system data
  WARNING: Make sure to back up your data
  Are you sure (N/y)?
  Activating staged installer...done.
  Reboot required to take effect.

Note: It is possible to combine the -i, -z and -a options all at once.
In addition the -f (force) option will suppress the yes/no question.

Example 4 - Stage and activate installer with ztp script.  No questions asked
-----------------------------------------------------------------------------

  cumulus@cumulus:~$ sudo onie-install -fa \
                          -i http://203.0.113.10/image-installer \
                          -z http://203.0.113.10/ztp-script
  Staging installer image... Adding ZTP script...done.
  Activating staged installer...done.
  Reboot required to take effect.

EOF
)

. /usr/share/cumulus/img/functions
cl_check_env root-priv



tmp_install_files=
mktemp_file()
{
    local tmp_var="$1"
    shift 1
    local t_file=$(mktemp "$@" --tmpdir=/tmp onie-install.XXXXXX)
    tmp_install_files="${tmp_install_files} $t_file"
    eval "${tmp_var}=\"${t_file}\""
}

# Clean up on exit function.  This function is only intended to be
# called by the signal trap handler.  It should not be called
# directly.
clean_up()
{
    /bin/rm -rf $tmp_install_files
}
# clean up when process exits or is killed gracefully
trap clean_up EXIT TERM

download_url()
{
    local url="$1"
    local dest="$2"

    # Sanitize the file URL
    local url_type=${url%%:*}
    local url_path=${url##*://}

    if [ "$url_type" = "$url" ] && [ "$url_path" = "$url" ] ; then
        # url is for a local file.  could be absolute or relative
        if [ -r "$url" ] ; then
            url=$(realpath $url)
        else
            log_failure_msg "$url is not a valid local file path."
            exit 1
        fi
    fi

    mktemp_file tmp_download_dir -d
    tmp_file="${tmp_download_dir}/$(basename $url_path)"

    case $url_type in
        http | ftp)
            _log_msg "${log_pre}Downloading URL: $url\n"
            if [ "$dry_run" = "n" ] ; then
                curl -# -f -o "$tmp_file" "$url" || {
                    log_failure_msg "Could not download file: $url."
                    return 1
                }
            fi
            log_success_msg "${log_pre}HTTP download complete."
            ;;
        https)
            _log_msg "${log_pre}Downloading URL: $url\n"
            if [ "$dry_run" = "n" ] ; then
                curl -# -f --insecure -o "$tmp_file" "$url" || {
                    log_failure_msg "Could not download file: $url."
                    return 1
                }
            fi
            log_success_msg "${log_pre}HTTPS download complete."
            ;;
        tftp)
            log_failure_msg "TFTP download is not supported."
            _log_msg "Use http or scp download instead.\n"
            exit 1
            ;;
        scp)
            _log_msg "${log_pre}Downloading URL: $url\n"
            if [ "$dry_run" = "n" ] ; then
                server=${url_path%%/*}
                # split on the /, but don't forget to put it back
                path=/${url_path#*/}
                scp ${server}:$path $tmp_file || {
                    log_failure_msg "Could not download file: $url."
                    return 1
                }
            fi
            log_success_msg "${log_pre}SCP download complete."
            ;;
        *)
            # assume it's a local path file
            _log_msg "${log_pre}Using file: $url\n"
            if [ -r "$url" ] ; then
                if [ -s "$url" ] ; then
                    if [ "$dry_run" = "n" ] ; then
                        cp "$url" "$tmp_file"
                    fi
                else
                    log_failure_msg "Requested file has zero size: $url"
                    return 1
                fi
            fi
            ;;
    esac

    if [ "$dry_run" = "n" ] ; then
        mv $tmp_file "$dest"
    fi
    rm -rf $tmp_download_dir

    return 0
}

status=y
install=n
ztp=n
cancel=n
activate=n
force=n
dry_run=n
log_pre=
cl_add_arg "scafnz:i:"
while getopts "${cl_args}" a ; do
    case $a in
        s)
            status=y
            ;;
        i)
            install=y
            url="$OPTARG"
            status=n
            ;;
        z)
            ztp=y
            ztp_url="$OPTARG"
            status=n
            ;;
        c)
            cancel=y
            status=n
            ;;
        a)
            activate=y
            status=n
            ;;
        f)
            force=y
            status=n
            ;;
        n)
            dry_run=y
            log_pre="Dry run: "
            printf "${log_pre}No changes will be made.\n"
            ;;
    esac
done

cl_get_args "$@"
shift $cl_arg_shift

stage_dir="/var/lib/cumulus/installer"
# Information kept in the stage directory:
# - onie-installer: the acutal image
# - status: human readable status info
# - launch: signal to initramfs

install_file="${stage_dir}/onie-installer"
install_status_file="${stage_dir}/status.install"
ztp_file="${install_file}.ztp"
ztp_status_file="${stage_dir}/status.ztp"
launch_file="${stage_dir}/launch"

status_info()
{
    if [ -r "$install_status_file" ] ; then
        _log_msg "Found installer staged for installation.\n"
        cat "$install_status_file"
    else
        _log_msg "No staged installer found.\n"
    fi
    if [ -r "$ztp_status_file" ] ; then
        _log_msg "Found zero touch provisioning (ZTP) script staged for installation.\n"
        cat "$ztp_status_file"
    else
        _log_msg "No staged zero touch provisioning (ZTP) script found.\n"
    fi
    if [ -r "$launch_file" ] ; then
        _log_msg "**************************************************\n"
        _log_msg "* Installer activation requested at next reboot. *\n"
        _log_msg "**************************************************\n"
    fi
}

progname_name="$(basename $0)"

{
    # Ensure only one instance is running
    flock -n 9 || {
        log_failure_msg "Another instance of '$progname_name' is already running..."
        exit 1
    }

    # Sanity check a few things
    if [ "$install" = "y" -a "$cancel" = "y" ] || \
           [ "$activate" = "y" -a "$cancel" = "y" ] || \
           [ "$ztp" = "y" -a "$cancel" = "y" ] ; then
        log_failure_msg "The cancel operation (-c) cannot be combined with other operations."
        exit 1
    fi

    if [ "$ztp" = "y" -a "$install" = "n" ] ; then
        log_failure_msg "The ztp URL option (-z) must be specified concurrently with the -i URL option."
        exit 1
    fi

    if [ "$status" = "y" ] ; then
        status_info
        exit 0
    fi

    if [ "$cancel" = "y" ] ; then
        if [ -d "$stage_dir" ] ; then
            if [ "$force" != "y" ] ; then
                printf "${log_pre}WARNING: Removing staged installer requested.\n"
                cl_prompt_yes_no "Are you sure" || exit 0
            fi
            log_begin_msg "${log_pre}Removing staged installer"
            [ "$dry_run" = "n" ] && rm -rf $stage_dir
            log_end_msg
        else
            _log_msg "No installer staged.  Nothing to do.\n"
        fi
        exit 0
    fi

    if [ "$install" = "y" ] ; then
        if [ "$dry_run" = "n" ] ; then
            rm -rf "$stage_dir"
            mkdir -p "$stage_dir"
        fi

        mktemp_file tmp_install_file

        _log_msg "Fetching installer: $url\n"
        download_url "$url" "$tmp_install_file" || {
            log_failure_msg "Unable to find installer URL: $url\n"
            exit 1
        }

        if [ "$dry_run" = "n" ] ; then
            chmod +x "$tmp_install_file"
            if grep -q CL_INSTALLER "$tmp_install_file" ; then
                mktemp_file tmp_info_file
                # Looks like a Cumulus installer - dump more info
                if "$tmp_install_file" info > $tmp_info_file ; then
                    # Verify switch architecture matches the current system
                    [ -r /etc/image-release ] && . /etc/image-release
                    image_switch_arch=$(egrep '^Switch-Architecture:' "$tmp_info_file" | awk '{ print $2 }')
                    if [ "$IMAGE_SWITCH_ARCHITECTURE" != "$image_switch_arch" ] ; then
                        log_failure_msg "Specified installer image does not match current switch architecture"
                        log_failure_msg "Current running switch architecture: $IMAGE_SWITCH_ARCHITECTURE"
                        log_failure_msg "Installer switch architecture      : $image_switch_arch"
                        log_failure_msg "Installer URL: $url"
                        exit 1
                    fi
                else if [ "$dry_run" = "n" ] ; then
                         log_failure_msg "$url is not a valid Cumulus installer image"
                         exit 1
                     fi
                fi
            else
                log_failure_msg "Specified image does not look like a Cumulus installer"
                log_failure_msg "Installer URL: $url"
                exit 1
            fi
        fi

        if [ "$ztp" = "y" ] ; then
            mktemp_file tmp_ztp_file
            _log_msg "Fetching ztp script: $ztp_url\n"
            download_url "$ztp_url" "$tmp_ztp_file" || {
                log_failure_msg "Unable to find ztp script URL: $ztp_url\n"
                exit 1
            }
            if [ "$dry_run" = "n" ] ; then
                # Validate ztp script
                if ! grep -q CUMULUS-AUTOPROVISIONING $tmp_ztp_file ; then
                    log_failure_msg "Specified ztp script missing marker 'CUMULUS-AUTOPROVISIONING'"
                    log_failure_msg "Invalid ztp script: $ztp_url"
                    exit 1
                fi
            fi
        fi

        log_begin_msg "${log_pre}Staging installer image"
        if [ "$dry_run" = "n" ] ; then
            mv "$tmp_install_file" "$install_file"
            echo "Installer-URL: $url" > "$install_status_file"
            [ -r "$tmp_info_file" ] && cat $tmp_info_file >> "$install_status_file"
            if [ "$ztp" = "y" ] ; then
                log_begin_msg " Adding ZTP script"
                mv "$tmp_ztp_file" "$ztp_file"
                echo "ZTP-URL: $ztp_url" > "$ztp_status_file"
                chmod +x "$ztp_file"
            fi
        fi
        log_end_msg
    fi

    if [ "$activate" = "y" ] ; then
        if [ -r "$install_file" ] ; then
            [ "$dry_run" = "n" ] && rm -f "$launch_file"
            if [ "$force" != "y" ] ; then
                printf "WARNING:\n"
                printf "WARNING: Activating staged installer requested.\n"
                printf "WARNING: This action will wipe out all system data.\n"
                printf "WARNING: Make sure to back up your data.\n"
                printf "WARNING:\n"
                cl_prompt_yes_no "Are you sure" || exit 0
            fi
            log_begin_msg "${log_pre}Activating staged installer"
            [ "$dry_run" = "n" ] && touch "$launch_file"
            log_end_msg
            if [ "$dry_run" = "n" ] ; then
                _log_msg "Reboot required to take effect.\n"
            fi
            exit 0
        else
            log_failure_msg "Activate operation requested, but no staged installer found."
            _log_msg "An installer must first be staged with the -i option.\n"
            exit 1
        fi
    fi

} 9> "/var/lock/${progname_name}.lock"
exit 0
