ec2-run

#!/bin/bash

if [ $# -lt 4 ]; then
	echo "Usage: [WATCH=3] ec2-run.sh environment purpose ssh-user [-f script.sh | command] [arg1] [arg2] [...] [argN]" 1>&2
	echo "Specifying a WATCH environment variable will loop with \$WATCH interval" 1>&2
	exit 1
fi

ENVIRONMENT="$1"
PURPOSE="$2"
USR="$3"
shift 3
## all the other arguments are passed as they are

## if first argument is '-f' then will use a local file as payload
if [[ "$1" = "-f" ]]; then
	SCRIPT="$2"
	shift 2
	SNAME="$(basename $SCRIPT)"
	## create a random prefix for the script
	RID=$RANDOM
fi

function run_script_on_all(){
	## run script on each host
	for HOST in $HOSTS; do
		echo -n "ssh -q '${USR}@$HOST' '/tmp/${RID}-${SNAME}' "
		echo "$@"
	done | coshell --deinterlace
}

function run_inline_on_all(){
	## run script on each host
	for HOST in $HOSTS; do
		ssh -q "${USR}@$HOST" "$@"
	done
}

set -e

echo -n "Getting matching EC2 hosts..."
HOSTS=$(ec2-find "$ENVIRONMENT" "$PURPOSE" | awk '{ print $NF }')
echo "OK"

if [ ! -z "$SCRIPT" ]; then
	## a serial operation will allow for SSH key ACKs
	for HOST in $HOSTS; do
		ssh "${USR}@$HOST" echo "$HOST: OK"
	done
	## first copy script in each host
	for HOST in $HOSTS; do
		echo "scp -q '$SCRIPT' '${USR}@$HOST:/tmp/${RID}-${SNAME}'"
	done | coshell --deinterlace
fi

set +e

if [ -z "$WATCH" ]; then
	if [ -z "$SCRIPT" ]; then
		run_inline_on_all "$@"
	else
		run_script_on_all "$@"
	fi
	RV=$?
else
	TMPF=$(mktemp)
	function remove_tmp() {
		rm "$TMPF"
	}
	trap remove_tmp EXIT

	while true; do
		echo "Every ${WATCH}s: $SCRIPT $@" > "$TMPF"
		if [ -z "$SCRIPT" ]; then
			run_inline_on_all "$@" 2>&1 >> "$TMPF"
			RV=$?
		else
			run_script_on_all "$@" 2>&1 >> "$TMPF"
			RV=$?
		fi
		echo "--- commands run finished at $(date) (exit code = $RV)" >> "$TMPF"
		clear
		cat "$TMPF"
		if [ ! $RV -eq 0 ]; then
			break
		fi
		sleep $WATCH
	done
fi

if [ ! -z "$SCRIPT" ]; then
	## remove temporary script
	for HOST in $HOSTS; do
		echo ssh -q "${USR}@$HOST" rm "/tmp/${RID}-${SNAME}"
	done | coshell --deinterlace
fi

exit $RV