#!/usr/bin/env bash

version="0.2"

# SSH Repeat
# version 0.2
#
# Copyright 2006 M&J Leonard Consulting Ltd.
#
# Not intended for distribution.

# The file .ssh-repeat can be used to specify options such as the EDITOR
# variable, so load it if it exists.
if [[ -e ~/.ssh-repeat ]]; then
    source ~/.ssh-repeat
fi

# Print a short usage message.
function print_usage () {
    echo "ssh-repeat version ${version}" 1>&2
    echo 1>&2
    echo "usage: ssh-repeat [-hp] [-o name] -[e|c cmd|f name]" 1>&2
    echo "    [-k keyfile] -[l name|g] [host1 host2 ...]" 1>&2

    exit 1
}

# Print a detailed usage message.
function print_help () {
    echo "ssh-repeat version ${version}" 1>&2
    echo 1>&2
    echo "usage: ssh-repeat [-hp] [-o name] -[e|c cmd|f name]" 1>&2
    echo "    [-k feyfile] -[l name|g] [host1 host2 ...]" 1>&2
    echo 1>&2
    echo "-h      Print this help information." 1>&2
    echo "-p      Continue to other hosts if one returns an error." 1>&2
    echo "        If this is not specified, ssh-repeat will terminate" 1>&2
    echo "        with the same code as the remote script in the event of" 1>&2
    echo "        an error." 2>&1
    echo "-k      Specify a public key file for use in authenticating" 1>&2
    echo "        with the server" 1>&2
    echo 1>&2
    echo "-c      Specify a command at the command line. Executed with" 1>&2
    echo "        /bin/sh on the target machine." 1>&2
    echo "-f      Specify a file name for a script. This will make a" 1>&2
    echo "        a temporary copy to avoid problems if the original is" 1>&2
    echo "        edited mid-run." 1>&2
    echo "-e      Interactively edit a script. When used alone, this will" 1>&2
    echo "        edit an empty file. When used in conjunction with -c or" 1>&2
    echo "        -f, this will edit a temporary copy of the script (no" 1>&2
    echo "        changes are made to the original)." 1>&2
    echo 1>&2
    echo "-l      Specifies a file name containing the list of hosts." 1>&2
    echo "        One per line, lines beginning with #'s are comments." 1>&2
    echo "-g      Gather the list of hosts from the script file (must be" 1>&2
    echo "        used in conjunction with -e or -f) from lines of the" 1>&2
    echo "        form '# SR:hostname'." 1>&2
    echo 1>&2
    echo "-o      Specify a bash script file to load before execution" 1>&2
    echo "        begins. This can be used to specify options such as the" 1>&2
    echo "        EDITOR variable." 1>&2
    echo 1>&2
    echo "NOTE:   Both scripts and host list files support includes" 1>&2
    echo "        from lines of the form '# SRI:filename'. Recursive" 1>&2
    echo "        includes are supported." 1>&2
    echo "NOTE:   Usernames can be specified in .ssh/config, or in the" 1>&2
    echo "        hostname specifications by prepending username@ to the" 1>&2
    echo "        server name." 1>&2

    exit 1
}

# Create a tempfile, and set the global command_file to its path.
function get_temp () {
    command_file="$( mktemp )"

    # terminate on error
    if [[ $? != 0 ]]; then
	echo -n ${0}: "couldn't get tempfile. " 1>&2
	echo "Check your TMPDIR variable." 1>&2
	exit 1
    fi
}

# Scans a file for include file lines, then runs the includefile routines on
# them.
function include_includes () {
    if [[ "${2}" ]]; then
	pushd "${2}" > /dev/null
    else
	pushd "$( dirname "${1}" )" > /dev/null
    fi

    for i in $( grep "^# SRI:" "${1}" | \
	sed -e "s,^# SRI:,," -e "s,^ *,," ); do
	echo "SR: including: ${i}" 2>&1
	file2hosts "${i}"
    done
    
    popd > /dev/null 2>&1
}

# Imports hosts from a file, adds them to the hosts array.
function file2hosts () {
    for i in $( grep -v "^#" "${1}" ); do
	hosts[${hosts_i}]="${i}"
	(( hosts_i++ ))
    done

    include_includes "${1}"
}

# Imports hosts from "# SR:" lines in a script.
function comments2hosts () {
    for i in $( grep "^# SR:" "${1}" | \
	sed -e "s,^# SR:,," -e "s,^ *,," ); do
      hosts[${hosts_i}]="${i}"
      (( hosts_i++ ))
    done

    include_includes "${1}"  "${basedir}"
}

# Deletes the script temp file so we don't clutter up /tmp.
function cleanup () {
    rm -f "${command_file}"
}

# Terminate cleanly, with no error.
function cleanexit () {
    cleanup
    exit 0
}

# Terminate with error and error message, cleanly.
function cleanexiterror () {
    echo "${myname}:" "${1}" 1>&2
    cleanup
    exit 1
}

# we make tempfiles, clean them up if someone wants to exit
trap cleanexit SIGINT

# so tempfiles aren't readable by others
umask 0077

# the only delimiter I care about is newlines
IFS=$'\n'

# verify we have an editor we can use
if [[ ! "${EDITOR}" ]]; then
    if [[ -x /usr/bin/vi ]]; then
	EDITOR="/usr/bin/vi"
    else
	echo ${0}: "couldn't find an editor. Check your EDITOR variable." 1>&2
	exit 1
    fi
fi

# I like to pretend I have to declare variables.
optstring="hpek:c:f:l:g"
myname="${0}"
opt=""
persist=""
declare -a hosts
hosts_i="0"
declare -a includes
includes_i="0"
hosts_gather=""
command_file=""
basedir=""
edit_script=""
key_file=""

# parse command line options
while getopts "${optstring}" OPT; do
    case $OPT in
	h)
	    print_help;;
	p)
	    persist="y";;
	k)
	    key_file="${OPTARG}";;
	e)
	    # remember the cwd for relative include paths
	    if [[ ! "${command_file}" ]]; then
		basedir="$( pwd )"
	    fi

	    # editing occurs later, once the source (if any) of the original
	    # is known
	    edit_script="y";;
	c)
	    if [[ "${command_file}" ]]; then
		cleanexiterror "Only one command or script can be used."
	    fi

	    get_temp

	    # create a script file from the command line specified
	    echo "#"\!"/bin/sh" >> "${command_file}"
	    echo "${OPTARG}" >> "${command_file}"

	    # remember the cwd for relative include paths
	    basedir="$( pwd )";;
	f)
	    if [[ "${command_file}" ]]; then
		cleanexiterror "Only one command or script can be used."
	    fi

	    if [[ ! -e "${OPTARG}" ]]; then
		cleanexiterror "The specified file does not exist."
	    fi

	    get_temp
	    # make a temporary copy of the command file
	    cp -f "${OPTARG}" "${command_file}"

	    # remember the directory of the script (for relative include paths)
	    basedir="$( dirname "${OPTARG}" )";;
	l)
	    if [[ ! -e "${OPTARG}" ]]; then
		cleanexiterror "The specified file does not exist."
	    fi

	    file2hosts "${OPTARG}";;
	g)
	    hosts_gather="y";;
    esac
done

# edit the script if that has been specified
if [[ "${edit_script}" ]]; then
    if [[ ! "${command_file}" ]]; then
	get_temp
    fi

    ${EDITOR} ${command_file}
fi

# if the command line parser doesn't know what the hell is going on, we've
# encountered some hosts at the end of the command line
if [[ $OPT == "?" && $OPTIND -le $# ]]; then
    for i in $( seq $OPTIND $# ); do
	hosts[${hosts_i}]="${!i}"
	(( hosts_i++ ))
    done
fi

if [[ ! "${command_file}" ]]; then
    cleanexiterror "No command was specified."
fi

if [[ "${hosts_gather}" ]]; then
    comments2hosts "${command_file}"
fi

if [[ ! "${hosts[@]}" ]]; then
    cleanexiterror "No hosts were specified."
fi

# a key file has been specified
if [[ "${key_file}" ]]; then
    ssh_opts="-i${key_file}"
    scp_opts="-i${key_file}"
fi

for i in "${hosts[@]}"; do
    # generate a tempfile on the destination
    destname="$( ssh ${ssh_opts} "${i}" mktemp /tmp/ssh-repeat.XXXXXX )"
    if [[ $? != 0 ]]; then
	cleanexit "Couldn't allocate a tempfile on ${i}. Exiting."
    fi
    
    echo "SR: ****************************************"
    echo "SR: ${i}"
    echo "SR: ****************************************"

    # upload the file
    scp ${scp_opts} -q "${command_file}" "${i}:${destname}"
    if [[ $? != 0 ]]; then
	cleanexit "Couldn't upload the script. Exiting."
    fi

    # execute the file, clean up
    ssh ${ssh_opts} "${i}" "chmod u+x ${destname}; ${destname}"
    code=$?
    ssh ${ssh_opts} "${i}" rm -f "${destname}"

    if [[ $code != 0 ]]; then
	echo "SR: ****************************************"
	echo "SR: ${i}: remote script terminated with error code ${code}"
	echo "SR: ****************************************"
	echo
	echo

	if [[ ! "${persist}" ]]; then
	    cleanup
	    exit ${code}
	fi
    else
	echo "SR: ****************************************"
	echo "SR: remote script finished"
	echo "SR: ****************************************"
	echo
	echo
    fi
done

cleanexit