#!/bin/bash

################################################################################
# blahp-over-ssh shim script
#
# command [options] [user@]hostname [options] command [remote arguments]
#
# Options:
#   see usage string 
#
# Arguments:
#   mandatory: hostname, command
#
# Remote arguments are passed along
#
# Assumes public/private key pair is already created and exists on host
#
# Exits with 0 on success, 1 if wrong arguments, < 0 if ssh-add failed, > 0 if ssh failed
################################################################################

#Defaults
REMOTE_CMD=""
REMOTE_HOSTNAME=""
REMOTE_USER=""
REMOTE_PORT=""
# Fix the home directory when spawned from a root process
unset HOME
export HOME=`echo ~`
# Needs '~' because resolved on remote site
REMOTE_GLITE="~/bosco/glite"
# We do this in case $HOME isn't set properly
PASSPHRASE_LOCATION=`echo ~/.bosco/.pass`
PRIVATE_KEY_LOCATION=`echo ~/.ssh/bosco_key.rsa`
# Set 'ssh -o BatchMode yes' by default so we fail if a password is requested
SSH_BATCHMODE=yes
BASH_LOGIN=yes
RVGAHP=no
RVGAHP_SOCKET=""

# Parse command line arguments 
PROG_NAME=$0
USAGE="Usage: $PROG_NAME [options] remote_hostname [options] REMOTE_CMD [remote arguments]\n \
$PROG_NAME [options] remote_hostname [remote options and arguments]\n \
Options: \n \
 --rgahp-user REMOTE_USER \tuser name on the remote host\n \
 --rgahp-key PRIVATE-KEY-PATH \tlocation of ssh private key (~/.ssh/bosco_key.rsa)\n \
 --rgahp-nokey  \t\tno ssh private key or key already enabled (same as empty rgahp-key)\n \
 --rgahp-pass PASSPHRASE \tlocation of passphrase protecting ssh private key (~/.bosco/.pass)\n \
 --rgahp-nopass \t\tno passphrase protecting ssh private key (same as empty rgahp-pass)\n \
 --rgahp-glite PATH \tpath to the directory of the script (~/bosco/glite)\n \
 --rgahp-script SCRIPT \tpath to script to start up blahp (PATH/bin/blahpd)\n \
 --rgahp-nobatchmode \t\tDisable SSH BatchMode \n \
 --rgahp-nologin \t\tDisable bash -l flag \n \
 --rvgahp-socket SOCKET_FILE \tUse the reverse gahp with the given socket file \n \
 --help, -h \t\t\tprint this\n \
 remote_hostname: [USER@]HOST[:PORT] same string that can be used to ssh to the host\n \
 remote arguments are passed to the REMOTE_CMD
 REMOTE_CMD can be expressed as argument or option
"

stop_ssh () {
    # Remove on-exit call of stop_ssh
    trap - EXIT

    # Shut down ssh-agent
    eval `ssh-agent -sk` 1>&2
}

while [ $# != 0 ] ; do
    case "$1" in
        -h | --help  ) echo -e "$USAGE"
            exit 0;;
        --rgahp-user  ) REMOTE_USER="$2"
            shift 2;;
        --rgahp-nokey  ) PRIVATE_KEY_LOCATION=""
            shift ;;
        --rgahp-key  ) PRIVATE_KEY_LOCATION="$2"
            shift ; shift  ;;
        --rgahp-pass  ) PASSPHRASE_LOCATION="$2"
            shift 2;;
        --rgahp-nopass  ) PASSPHRASE_LOCATION=""
            shift ;;
        --rgahp-glite ) REMOTE_GLITE="$2"
            shift 2;;
        --rgahp-script ) REMOTE_CMD="$2"
            shift 2;;
        --rgahp-nobatchmode ) SSH_BATCHMODE=no
            shift ;;
        --rgahp-nologin ) BASH_LOGIN=no
            shift ;;
        --rgahp-* ) echo "Unknown option: $1" 1>&2
            echo -e "$USAGE"
            exit 1;;
        --rvgahp-socket ) RVGAHP_SOCKET="$2"
            RVGAHP=yes
            shift 2;;
        -- ) shift
            if [ "$REMOTE_CMD" = "" -o "$1" = "condor_ft-gahp" ] ; then
                REMOTE_CMD="$1"
                shift
            fi
            break;;
        -* ) break;;
        * )
            if [ "$REMOTE_HOSTNAME" = "" ] ; then
                if  [[ $1 =~ ^(([A-Za-z_][A-Za-z0-9_-]*)@)?([A-Za-z0-9\.-]+)(:([0-9]*))?$ ]]; then
                    REMOTE_USER="${BASH_REMATCH[2]:-$REMOTE_USER}"
                    REMOTE_HOSTNAME="${BASH_REMATCH[3]:-$REMOTE_HOSTNAME}"
                    REMOTE_PORT="${BASH_REMATCH[5]:-$REMOTE_PORT}"
                fi
                shift
            elif [ "$REMOTE_CMD" = "" -o "$1" = "condor_ft-gahp" ] ; then
                REMOTE_CMD="$1"
                shift
            else
                break
            fi
            ;;
    esac
done

if [ "$REMOTE_HOSTNAME" = "" ] ; then
    echo "Missing hostname" 1>&2
    echo -e "$USAGE"
    exit 1
fi
if [ "$REMOTE_CMD" = "" ] ; then
    echo "Missing remote command" 1>&2
    echo -e "$USAGE"
    exit 1
fi

if [ "$RVGAHP" = "yes" ] ; then
    # TODO handle default RVGAHP_SOCKET
    exec $(condor_config_val SBIN)/rvgahp_client "${RVGAHP_SOCKET}" "${REMOTE_CMD}"
    echo "Failed to run rvgahp_client" 1>&2
    exit 1
fi

##### Handling authentication #####
# Start and init ssh agent if key file is specified


# if a ssh key is required, start up a ssh-agent and do ssh-add
if [ -n "$PRIVATE_KEY_LOCATION" -a -f "$PRIVATE_KEY_LOCATION" ] ; then
	# Run stop_ssh on shell exit
	trap stop_ssh EXIT

	# start the ssh-agent
	eval `ssh-agent -s` 1>&2

	# Call the external program to do ssh-add
        # If a passphrase is required pass it to the script
        if [ -n "$PASSPHRASE_LOCATION" ]; then
	    condor_ssh_start --key "$PRIVATE_KEY_LOCATION" --pass "$PASSPHRASE_LOCATION"
        else
	    condor_ssh_start --key "$PRIVATE_KEY_LOCATION" --nopass
        fi
	ADD_STATUS=$?

	# check if ssh-add failed
	if [ $ADD_STATUS != 0 ] ; then
		exit $ADD_STATUS
	fi
fi


##### Running remote command and cleanup #####

# remove hostname from arglist
shift

[[ $REMOTE_PORT != "" ]] && SSH_ARGS=(-p $REMOTE_PORT)
# Some HPCs require a PIN when specifying "-o BatchMode yes", so we should allow users
# https://opensciencegrid.atlassian.net/browse/SOFTWARE-4239
[[ $SSH_BATCHMODE == "yes" ]] && SSH_ARGS+=(-o "BatchMode yes")

[[ $BASH_LOGIN == "yes" ]] && BASH_ARGS+=(-l)

if [[ $REMOTE_USER != "" ]] ; then
    HOSTNAME_ARG="${REMOTE_USER}@${REMOTE_HOSTNAME}"
else
    HOSTNAME_ARG="${REMOTE_HOSTNAME}"
fi

if [[ $REMOTE_CMD == "blahpd" || $REMOTE_CMD == "batch_gahp" ]] ; then
    ssh "${SSH_ARGS[@]}" $HOSTNAME_ARG /bin/bash "${BASH_ARGS[@]}" -c "'BLAHPD_LOCATION=$REMOTE_GLITE GLITE_LOCATION=$REMOTE_GLITE $REMOTE_GLITE/bin/$REMOTE_CMD $*'"
    SSH_STATUS=$?
elif [ "${REMOTE_CMD}" = "condor_ft-gahp" ] ; then
    # We need to set up a tunnel from the remote machine for the file
    # transfer TCP connections. If we knew that both sides were running
    # OpenSSH 5.2p1 or later, we could have ssh pick the port on the
    # remote end. But we don't, so we pick a random port and hope it's
    # not already in use.
    # We mimic the message that newer versions of OpenSSH print when
    # binding a dynamic port for tunneling. The gridmanager looks for
    # this message to know which port to tell the ft-gahp to use.
    # If the local OpenSSH is 4.4p1 or later (we check for 5.0 or later
    # for simplicity), then we can use ExitOnForwardFailure and try
    # several random ports in case we get unlucky on the first attempt.
    # We extract the IP and port on which the gridmanager can be
    # contacted from $CONDOR_INHERIT.
    if [ "$CONDOR_INHERIT" != "" ] ; then
        GRIDMANAGER_ADDRESS=`echo "$CONDOR_INHERIT" | sed 's/[^<]*<\([^?>]*\).*/\1/'`
    else
        GRIDMANAGER_ADDRESS="127.0.0.1:12345"
        echo "CONDOR_INHERIT not defined, using bogus value $GRIDMANAGER_ADDRESS for gridmanager address" 1>&2
    fi
    SSH_STATUS=255
    if [[ `ssh -V 2>&1` =~ ^OpenSSH_[5-9].* ]] ; then
        SSH_ARGS+=(-o ExitOnForwardFailure=yes)
        tries=3
    else
        tries=1
    fi
    while ((tries-- > 0 && SSH_STATUS == 255)) ; do
        let port=${RANDOM}+32768
        ssh "${SSH_ARGS[@]}" -R $port:$GRIDMANAGER_ADDRESS $HOSTNAME_ARG /bin/bash "${BASH_ARGS[@]}" -c "'echo Allocated port $port for remote forward to 1>&2 ; CONDOR_CONFIG=$REMOTE_GLITE/etc/condor_config.ft-gahp $REMOTE_GLITE/bin/condor_ft-gahp -f $*'"
        SSH_STATUS=$?
    done
else
    echo "Unknown remote command (${REMOTE_CMD})" 1>&2
    SSH_STATUS=1
fi

exit $SSH_STATUS
