tmux-cssh

tmux-cssh.zsh

#!/usr/bin/env zsh
# tmux cluster ssh + knife support
# Daniel Kertesz <[email protected]>
#
# Add basic hosts completion to zsh by running:
# compdef _hosts tmux-cssh

set -e
autoload colors; colors

function fatal() {
    echo "$fg_bold[red]ERROR${reset_color}: $@" >&2
    exit 1
}

function warn() {
    echo "$fg_bold[yellow]WARNING${reset_color}: $@" >&2
}

function phelp() {
    printf "%-15s\t%s\n" "$1" "$2"
}

# Sort first by data-center, then by Pod, then by server id
function sort_hosts() {
    sort -t '.' -k 3,3 -k 2,2 -k 1,1
}

# 12 servers can be viewed in a 3x4 split window
HOSTLIMIT=12

user_query=""
query=()
ssh_options=()
do_log=0
do_force=0
do_dry_run=0
do_bg=0
pods=()
session_names=()

while getopts ":r:p:q:l:LPbvfnh" opt; do
    case $opt in
        L)
            do_log=1
            ;;
        r)
            query+="role:hg_$OPTARG"
            session_names+="$OPTARG"
            ;;
        p)
            pods+="$OPTARG"
            session_names+="pod$OPTARG"
            ;;
        P)
            query+='chef_environment:*production*'
            ;;
        b)
            do_bg=1
            ;;
        q)
            user_query="$OPTARG"
            ;;
        l)
            ssh_options+="-l$OPTARG"
            ;;
        n)
            do_dry_run=1
            ;;
        v)
            ssh_options+="-v"
            ;;
        f)
            do_force=1
            ;;
        n)
            do_dry_run=1
            ;;
        h)
            printf "Usage: $(basename $0) [-nLPbfvh] [-p pod] [-q query] [-r role] [-l user] [host ...]\n"
            printf "NOTE: hostnames will be read from stdin when no arguments are specified.\n\n"
            phelp "-n" "Dry run: Print the results of 'knife search', do not execute ssh"
            phelp "-l <user>" "Specifies the user to log in as on the remote machines"
            phelp "-L" "Capture each pane output to a file"
            phelp "-p <pod>" "Search by Pod number (can be specified multiple times)"
            phelp "-P" "Limit the search to '*production*' nodes"
            phelp "-b" "Do not switch to the created tmux session"
            phelp "-q <query>" "Specify a custom knife search query"
            phelp "-r <role>" "Search by hostgroup role (without the 'hg_' prefix)"
            phelp "-f" "Suppress the error when more than $HOSTLIMIT hosts are matched"
            phelp "-v" "Run a verbose ssh session"
            phelp "-h" "Show this help"
            exit
            ;;
        \?)
            fatal "Invalid option: -$OPTARG"
            ;;
        :)
            fatal "Option -$OPTARG requires an argument"
            ;;
    esac
done

shift $((OPTIND-1))

hosts=()

# Check for conflicting command line switches
if [[ $#query > 0 && ! -z $user_query ]]; then
    fatal "Can't specify a custom query (-q) together with filtering options (-p -P -r)"
fi

# create filters for the selected POD(s)
if [[ $#pods = 1 ]]; then
    query+="zendesk_config_pod:$pods[1]"
elif (( $#pods > 1 )); then
    # to understand the following line search for ${^spec} in man zshexpn.
    # Creates an array that contains (zendesk_config_pod:X zendesk_config_pod:Y ...).
    local p=("zendesk_config_pod:"${^pods})
    # then join the array with a OR and wrap the results in a subquery for knife.
    query+="(${(j: OR :)p})"
fi

# Build our knife search query
# NOTE: do not redirect knife stderr, we need to see errors.
if [[ $#query > 0 ]]; then
    # using knife search
    hosts=($(knife search node -i "${(j: AND :)query}"))
elif [[ ! -z $user_query ]]; then
    # using custom user knife search
    hosts=($(knife search node -i "$user_query"))
elif [[ $# == 0 ]]; then
    # reading hosts from stdin
    warn "Reading server list from stdin"
    while read line; do
        hosts+=$line
    done
else
    # reading hosts from command line
    hosts+=($@)
fi

# sort the hosts array to have nice ordered tmux windows later
IFS=$'\n' hosts=($(sort_hosts <<<"${hosts[*]}"))
unset IFS

# Print and exit if the user requested a dry-run
if [[ $do_dry_run = 1 ]]; then
    print ${(j:\n:)hosts}
    exit 0
fi

# Safety checks
if [[ $#hosts = 0 ]]; then
    fatal "No hosts matched your query or you gave me an empty list"
elif (( $#hosts > $HOSTLIMIT )) && [[ $do_force = 0 ]]; then
    echo -n "Too many hosts matched ($#hosts); are you sure you want to continue? [y/N]"
    if ! read -sq; then
        exit 1
    fi
    echo
fi

SESSION="cssh-$$"
if (( $#session_names > 0 )); then
    SESSION="$SESSION-${(j:-:)session_names}"
fi
WINDOW="${SESSION}:1"

now=$(date +"%d-%m-%Y-%H:%M")
logname="$PWD/tmux-cssh.${now}.#P.HOST.log"

i=1

tmux new-session -d -s $SESSION -P -F "tmux session: #{session_name}" \
     "set_pane_title $hosts[$i] ; ssh $ssh_options[@] $hosts[$i]"
[[ $do_log = 1 ]] && tmux pipe-pane -t ${WINDOW}.$i "cat >> ${logname/HOST/$hosts[$i]}"
i=$((i+1))

while [[ $i -le $#hosts ]]; do
    tmux split-window -t $WINDOW "set_pane_title $hosts[$i] ; ssh $ssh_options[@] $hosts[$i]"
    tmux select-layout -t $WINDOW tiled
    [[ $do_log = 1 ]] && tmux pipe-pane -t ${WINDOW}.$i "cat >> ${logname/HOST/$hosts[$i]}"
    i=$((i+1))
done

tmux setw -t $WINDOW status on
tmux setw -t $WINDOW synchronize-panes on
tmux set-hook -t $SESSION pane-exited "select-layout tiled"

if [[ $do_bg = 0 ]]; then
    tmux switch-client -t $SESSION
fi