tmc

#!/bin/sh

# Copyright 2015 David Scholberg <recombinant.vector@gmail.com>

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# define newline var, per http://stackoverflow.com/a/9403076
NL='
'

usage() {
	cat << EOF
Usage: $(basename "$0") [OPTIONS] [CLUSTERNAME]

Options:
    -h              Display help message and quit
    -d              Dump space-separated list of hosts in CLUSTERNAME
    -c CLUSTERLINE  Create custom cluster
                        CLUSTERLINE is treated as a clusters config line.
                        This option conflicts with passing CLUSTERNAME.
                        The name used for the cluster must not exist in the
                            clusters config.
EOF
}

# Add cluster hosts to global $HOSTS
# param 1 should be cluster name
add_cluster_hosts() {
	local CLUSTER="$1"

	# check for cluster in config
	local CLUSTER_LINE="$(echo "$CONF_LINES" | grep -E "^$CLUSTER ")"
	if [ -z "$CLUSTER_LINE" ]; then
		HOSTS="$HOSTS $CLUSTER"
	else
		SEEN_CLUSTERS="$SEEN_CLUSTERS $CLUSTER"

		# get hosts from cluster line
		local CLUSTER_LINE_HOSTS="$(echo "$CLUSTER_LINE" | cut -f 2- -d ' ')"

		for HOST in $CLUSTER_LINE_HOSTS; do
			if [ -z "$(echo "$HOSTS" | grep -E " $HOST( |$)")" -a -z "$(echo "$SEEN_CLUSTERS" | grep -E " $HOST( |$)")" ]; then
				add_cluster_hosts "$HOST"
			fi
		done
	fi

}

# check if tmux is in PATH
if [ -z "$(which tmux)" ]; then
	echo "error: tmux not found in path: $PATH" 1>&2
	exit 1
fi

CONF="$HOME/.clusterssh/clusters"

# check for conf file
if [ ! -f "$CONF" ]; then
	echo "error: config $CONF not found" 1>&2
	exit 2
fi

DUMP_HOSTS=""
CUSTOM_CLUSTER_LINE=""

# parse args
while getopts :hdc: OPT; do
	case $OPT in
		h)
			usage
			exit
			;;
		d)
			DUMP_HOSTS="true"
			;;
		c)
			CUSTOM_CLUSTER_LINE="$OPTARG"
			;;
		\?)
			echo "error: invalid option -$OPTARG" 1>&2
			usage 1>&2
			exit 3
			;;
		:)
			echo "error: missing argument for option -$OPTARG" 1>&2
			usage 1>&2
			exit 3
			;;
	esac
done

shift $((OPTIND-1))

CONF_LINES="$(grep -Ev '(^#|^$)' "$CONF")"

if [ -n "$CUSTOM_CLUSTER_LINE" ]; then
	# check for first param
	if [ -n "$1" ]; then
		echo "error: passing CLUSTERNAME conflicts with -c option"
		usage 1>&2
		exit 3
	fi

	# get cluster name
	CLUSTER="$(echo "$CUSTOM_CLUSTER_LINE" | awk '{print $1}')"

	# check conf for existing cluster of the same name
	if [ -n "$(echo "$CONF_LINES" | grep -E "^$CLUSTER ")" ]; then
		echo "error: cluster $CLUSTER specified with -c option exists in config $CONF"
		usage 1>&2
		exit 3
	fi

	# add custom config line to configuration
	CONF_LINES="${CONF_LINES}${NL}${CUSTOM_CLUSTER_LINE}${NL}"
else
	# check for first param
	if [ -z "$1" ]; then
		usage 1>&2
		exit 3
	fi

	CLUSTER="$1"
fi

# check for cluster in config
CLUSTER_LINE="$(echo "$CONF_LINES" | grep -E "^$CLUSTER ")"
if [ -z "$CLUSTER_LINE" ]; then
	echo "error: cluster $CLUSTER not in config $CONF" 1>&2
	exit 3
fi

# check for existing tmux session of the specified cluster
if [ -n "$(tmux ls | grep "cluster-$CLUSTER:")" ]; then
	echo "error: tmux session cluster-$CLUSTER already exists" 1>&2
	exit 3
fi

# get hosts from cluster line
HOSTS=""
SEEN_CLUSTERS=""
add_cluster_hosts $CLUSTER

# remove leading space
HOSTS="$(echo "$HOSTS" | sed 's/^ //')"

if [ -n "$DUMP_HOSTS" ]; then
	echo "$HOSTS"
	exit
fi

# build tmux commands
# get first host
HOST="$(echo "$HOSTS" | awk '{print $1}')"
# remove first host
if [ "$(echo "$HOSTS" | grep -o ' ' | wc -l)" -gt 0 ]; then
	HOSTS="$(echo "$HOSTS" | cut -f 2- -d ' ')"
else
	HOSTS=""
fi

SHELL_CMD="ssh $HOST; [ \$? -eq 255 ] && (echo Press ENTER to close pane; read enter)"
TMUX_COMMANDS="new -s \"cluster-$CLUSTER\" \"$SHELL_CMD\"$NL"

for HOST in $HOSTS; do
	SHELL_CMD="ssh $HOST; [ \$? -eq 255 ] && (echo Press ENTER to close pane; read enter)"
	TMUX_COMMANDS="${TMUX_COMMANDS}splitw \"$SHELL_CMD\"$NL"
	TMUX_COMMANDS="${TMUX_COMMANDS}select-layout tiled$NL"
done

TMUX_COMMANDS="${TMUX_COMMANDS}set-window-option synchronize-panes on$NL"

TMPFILE="$(mktemp /tmp/tmux-cluster-XXXXXXXXXXXX)"

echo "$TMUX_COMMANDS" > "$TMPFILE"
tmux source-file "$TMPFILE"
rm "$TMPFILE"