#!/usr/bin/env bash
# rdo -- run a command across all hosts

. lib.bash || exit

usage() {
	echo "Usage: $progname [-H hosts] [-C path|-N|-P] [-SWq] <command>"
	echo
	echo "Run a command simultaneously on all hosts."
	echo
	echo_opt "-H hosts"	"list of hosts to connect to"
	echo_opt "-S"		"connect as superuser (root)"
	echo
	echo_opt "-C path"	"set working directory"
	echo_opt "-N"		"use local working directory via NFS"
	echo_opt "-P"		"don't load profile (non-Unix OS)"
	echo
	echo_opt "-W"		"output in order (disables parallel output)"
	echo_opt "-q"		"don't show hostnames or exit values"
	# XXX: We could still invoke the commands in parallel and only show
	# output in order, after all have returned. This might increase the
	# perceived delay, though, so it should be -WW or something.
}

wd=
noshell=0
nfs=0
hosts=
user=
wait=0
quiet=0

while getopts ":C:H:NPSWq" OPT; do
	case $OPT in
	C) wd=$OPTARG;;
	H) hosts=${OPTARG//,/ };;
	N) nfs=1;;
	P) noshell=1;;
	S) user="root@";;
	W) wait=1;;
	q) quiet=1;;
	*) lib:die_getopts;;
	esac
done; shift $((OPTIND-1))

hosts=$(rlisthosts "$hosts") || exit

if (( ! $# )); then
	die "command not specified"
elif (( $# == 1 )); then
	cmd=$1
else
	cmd=${*@Q}
fi

if (( nfs )); then
	if [[ ! -e /proc/fs/nfsd/versions ]]; then
		die "NFS service not running on $HOSTNAME"
	fi
	wd=${wd:-"."}
fi

if [[ $wd ]]; then
	# Allow "-C ." to be specified in a useful way (i.e. translate the
	# local working directory into its remote or NFS equivalent)
	if [[ $wd != /* ]]; then
		wd=$(realpath -s "$wd")
	fi
	if (( nfs )); then
		wd="$wd/"
		wd="/net/$HOSTNAME/${wd#/net/*/}"
		wd="${wd%/}"
	fi
fi

if [[ $wd ]]; then
	if (( noshell )); then
		wdcmd="cd ${wd@Q} && $cmd"
	else
		wdcmd="(cd ${wd@Q} && \$SHELL -l -c ${cmd@Q}) 2>&1"
	fi
else
	if (( noshell )); then
		wdcmd="$cmd"
	else
		wdcmd="\$SHELL -l -c ${cmd@Q} 2>&1"
	fi
fi

lck=$(mktemp /tmp/rdo.XXXXXXXX) || exit
trap 'wait; rm -f "$lck"' INT

for host in $hosts; do
	{
		case $host in
			# Assume -P -C "" for '!'-suffixed hosts
			# (allows them to be defined in rlisthosts)
			*!) host=${host%!};;
			*) cmd=$wdcmd;;
		esac
		out=$(mktemp /tmp/rdo.XXXXXXXX) || exit
		if (( quiet )); then
			ssh -n $user$host "$cmd" > "$out"
			flock "$lck" cat "$out"
		else
			(ssh -n $user$host "$cmd" || echo "=> $?") > "$out"
			flock "$lck" awk -v host=$host \
				'BEGIN {printf "%s:", host}
				{print "\t" $0}
				END {if (!NR) print}' < "$out"
		fi
		rm -f "$out"
	} &
	if (( wait )); then wait; fi
done

wait; rm -f "$lck"
