tmc

#!/bin/sh

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

# This file is part of tmux-cluster.
#
# tmux-cluster 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.
#
# tmux-cluster 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 tmux-cluster.  If not, see <http://www.gnu.org/licenses/>.

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

# error codes
ERR_DEPENDENCY=1
ERR_CONFIG=2
ERR_ARG=3

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

Options:
    -h              Display help message and quit.
    -d              Dump space-separated list of hosts in CLUSTERNAME.
                        Conflicts with -t option.
    -t              Dump generated list of tmux commands.
                        Conflicts with -d option.
    -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.
    -x EXCLUDES     Space-separated list of hosts to exclude from cluster.
                        The hosts in EXCLUDES will not be connected to.
    -w              Create cluster panes in a new window in the current session.
                        Only valid if used in an attached tmux session.
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

}

# get dimensions of current tmux window
get_current_tmux_window_dimesions() {
    echo "$(tmux list-windows | grep -E '[0-9]+: \w+\*' | sed -r 's/^.+\[([0-9]+x[0-9]+)\].+$/\1/')"
}

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

DUMP_HOSTS=""
DUMP_TMUX_CMDS=""
CUSTOM_CLUSTER_LINE=""
EXCLUDED_HOSTS=""
USE_EXISTING_SESSION=""

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

shift $((OPTIND-1))

# if -w option is used, make sure we're in tmux
if [ "$USE_EXISTING_SESSION" = "true" -a -z "$TMUX" ]; then
    MSG="error: -w option must be used in an attached tmux session"
    echo "$MSG" 1>&2
    usage 1>&2
    exit "$ERR_ARG"
fi

# check for conficting dump options
if [ -n "$DUMP_HOSTS" -a -n "$DUMP_TMUX_CMDS" ]; then
    echo "error: -d and -t options conflict" 1>&2
    usage 1>&2
    exit "$ERR_ARG"
fi

CONF="$HOME/.clusterssh/clusters"

# check for conf file or custom cluster option
if [ ! -f "$CONF" -a -z "$CUSTOM_CLUSTER_LINE" ]; then
    echo "error: either config $CONF must exist or -c option must be used" 1>&2
    usage 1>&2
    exit "$ERR_ARG"
fi

CONF_LINES=""

if [ -f "$CONF" ]; then
    CONF_LINES="$(grep -Ev '(^#|^$)' "$CONF")"
fi

if [ -n "$CUSTOM_CLUSTER_LINE" ]; then
    # check for first param
    if [ -n "$1" ]; then
        echo "error: passing CLUSTERNAME conflicts with -c option" 1>&2
        usage 1>&2
        exit "$ERR_ARG"
    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" 1>&2
        usage 1>&2
        exit "$ERR_ARG"
    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 "$ERR_ARG"
    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 "$ERR_ARG"
fi

SESSION_NAME=""
# set session name
if [ -z "$USE_EXISTING_SESSION" ]; then
    SESSION_NAME="cluster-$CLUSTER"

    # check for existing tmux session of the specified cluster
    if [ -z "$DUMP_HOSTS" -a -z "$DUMP_TMUX_CMDS" ]; then
        # TODO: update to use has-session tmux command
        if [ -n "$(tmux ls 2>/dev/null | grep "$SESSION_NAME:")" ]; then
            echo "error: tmux session $SESSION_NAME already exists" 1>&2
            exit "$ERR_ARG"
        fi
    fi
fi

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

# exclude hosts
# TODO: allow EXCLUDED_HOSTS to also contain cluster names
if [ -n "$EXCLUDED_HOSTS" ]; then
    for HOST in $EXCLUDED_HOSTS; do
        HOSTS="$(echo "$HOSTS" | sed -r "s/ $HOST( |$)/ /g")"
    done
fi

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

# dump hosts
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)"

# if we're running in tmux, get dimensions of the current window
# this is needed if tmc is run inside run-shell, since the default window size is 80x24
TMUX_WINDOW_XY_ARGS=""
if [ -n "$TMUX" ]; then
    TMUX_DIMENSIONS="$(get_current_tmux_window_dimesions)"
    TMUX_WINDOW_X="$(echo "$TMUX_DIMENSIONS" | cut -f 1 -d 'x')"
    TMUX_WINDOW_Y="$(echo "$TMUX_DIMENSIONS" | cut -f 2 -d 'x')"
    TMUX_WINDOW_XY_ARGS="-x $TMUX_WINDOW_X -y $TMUX_WINDOW_Y"
fi

TMUX_SESSION_S_ARG=""
TMUX_SESSION_T_ARG=""
if [ "$USE_EXISTING_SESSION" = "true" ]; then
    TMUX_CMDS="neww \"$SHELL_CMD\"$NL"
else
    TMUX_SESSION_S_ARG="-s \"$SESSION_NAME\""
    TMUX_SESSION_T_ARG="-t \"$SESSION_NAME\""
    TMUX_CMDS="new -d $TMUX_WINDOW_XY_ARGS $TMUX_SESSION_S_ARG \"$SHELL_CMD\"$NL"
fi

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

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

# if we're inside a tmux client, switch-client, else attach
if [ -n "$TMUX" ]; then
    # only switch client if we created a new session
    # if using an existing session, we're already in the new window
    if [ -z "$USE_EXISTING_SESSION" ]; then
        TMUX_CMDS="${TMUX_CMDS}switch-client $TMUX_SESSION_T_ARG$NL"
    fi
else
    TMUX_CMDS="${TMUX_CMDS}attach $TMUX_SESSION_T_ARG$NL"
fi

# fix issue with incorrect layout after call to switch-client
TMUX_CMDS="${TMUX_CMDS}select-layout $TMUX_SESSION_T_ARG tiled$NL"

# dump tmux commands
if [ -n "$DUMP_TMUX_CMDS" ]; then
    echo -n "$TMUX_CMDS"
    exit
fi

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

echo "$TMUX_CMDS" > "$TMPFILE"
tmux start-server \; source-file "$TMPFILE"
rm "$TMPFILE"