#!/bin/bash
#
# 2014/03/26 Gabriel Moreau <Gabriel Moreau(A)univ-grenoble-alpes.fr> - Initial release
#
# From http://hd-recording.at/dokuwiki/doku.php?id=linux:tmux#tssh

# Clean when Ctrl^C
trap '[ -n "${base_path}" -a -d "/tmp/${base_path}" ] && rm -rf "/tmp/${base_path}"; exit 4;' QUIT INT TERM

export PATH=/bin:/sbin:/usr/bin:/usr/sbin:/usr/local/bin
export LANG=C


function usage() {
   cat <<END_USAGE
NAME
   $(basename $0) - tmux cluster ssh

SYNOPSIS
   $0 [-w number] [-f] [-v] [-c remote_cmd] <host1> <host2> <clusterssh class>... <hostM>- <hostN>+

OPTIONS
   -w             windows to open (integer, default 16)
   -f             fast, no nmap scan to eliminate sleeping computer
   -v             verbose
   -c remote_cmd  launch the remote command on hosts and exit
   -h             help

DESCRIPTION
   tssh can be use to launch terminal on many computer in parallel with tmux
   multiplexer and ssh.
   The tmux windows is splitted automatically.
   If you need more computers on the same windows, you can zoom in and out
   under gnome terminal with Ctrl- or Ctrl+.
   This must be done before launching tssh.
   
   On the command line, you can put host, login@host, clusterssh class.
   A host or a class can be remove from the list with a dash append
   and force to be in this one with a plus append.
   Example with the cluster ssh config below:
   
    tssh all team- node005 laptop04+

   Is equivalent to:
 
    tssh srv-mail srv-dns srv-imap srv-web srv-proxy \\
      node001 node002 node003 node004 \\
      node101 node102 node103 node104 \\
      node005 laptop04

   The control command for tmux is Ctrl^b.
   You can switch from broadcast to a local machine with Ctrl^b Ctrl^b
   and move between machine with Ctrl^b ArrowKey.

DEPENDS
   On Debian, you need the package

    apt-get install tmux ncurses-bin wamerican nmap

   wamerican (or wfrench...) is used to choose a random word in the file /usr/share/dict/words
   for each new tmux session.

   ncurses-bin is required for the tput command
   to automatically split your terminal into several small panels.
   nmap is only used for dynamic DNS domain and dynamic scan.
   This is not mandatory for general use.

   By default, tssh use tput to know the number of columns and lines of your terminal.
   It takes 10 lines and 40 columns for each windows by default.
   If tput is not installed, the default is 16 windows...

CONFIGURATION
   The clusterssh config file ~/.csshrc is a key values file.
   The "clusters" is mandatory for clusterssh (not tssh) and define the other keys.
   Values could be computer list or other key...
   
    clusters = all server s1 s2 s3 node n1 n2 team switch
    all = server node team
    server = s1 s2
    node = n1 n2
    s1 = srv-mail srv-dns srv-imap
    s2 = srv-web srv-proxy
    n1 = node001 node002 node003 node004
    n2 = node101 node102 node103 node104
    team = pc01 pc06 laptop04 laptop05 laptop09
    switch = root@switch01 root@switch05 root@switch17

   The tssh config file (~/.tsshrc) can be use change the default parameters.
   
    #export split_number=16
    export dyn_domain='mycompagny.local'
    #export fast='yes'
    #export verbose='yes'

AUTHOR
   Gabriel Moreau

COPYRIGHT
   Copyright (C) 2014-2018, LEGI UMR 5519 / CNRS UGA G-INP, Grenoble, France
   Licence : GNU GPL version 2 or later
END_USAGE
   }

export remote_command=''
export split_number=16
if which tput > /dev/null
then
   export split_number=$(( ($(tput  lines)/ 10) * ($(tput  cols)/ 40) ))
fi

export dyn_domain=''
if [ -e "${HOME}/.tsshrc" ]
then
   . "${HOME}/.tsshrc"
fi

# get options
if [ $# -eq 0 ]; then usage; exit 1; fi 
while getopts "w:c:fvh" options
do
   case ${options} in
      w)
         if echo ${OPTARG} | egrep -q '^[[:digit:]]+$' && [ ${OPTARG} -gt 0 ]
         then
            export split_number=${OPTARG}
         else
            usage
            exit 2
         fi
         ;;
      c)
         if echo ${OPTARG} | egrep -q '[[:alpha:]]'
         then
            export remote_command=${OPTARG}
         else
            usage
            exit 2
         fi
         ;;
      f)
         export fast='yes'
         ;;
      v)
         export verbose='yes'
         ;;
      h|*)
         usage
         exit 3
         ;;
   esac
done
shift $((OPTIND - 1))
[[ $1 = "--" ]] && shift

cd /tmp/
export base_path=$(mktemp -d tssh.XXXXXX)
touch "/tmp/${base_path}/master"
touch "/tmp/${base_path}/master-"
touch "/tmp/${base_path}/master+"
touch "/tmp/${base_path}/master--"
touch "/tmp/${base_path}/master++"

get_host_list () {
   local cluster
   local default_mode=''

   # set local mode
   if echo $1 | grep -- ^--mode=
   then
      default_mode=$(echo $1 | grep -- ^--mode= | cut -f 2 -d '=')
      shift
   fi

   for host in $*
   do
      local mode=${default_mode}
      local last_char="${host: -1}"
      if [ "${last_char}" == "-" -o "${last_char}" == "+" ]
      then
         mode="${last_char}"
         host="${host:0:${#host}-1}"
      fi

      # short host without login part if any
      local justhost=${host#*@}
      
      cluster=$(grep "^${justhost}\b" ${HOME}/.csshrc | cut -f 2 -d '=' | sed -e 's/^[[:space:]]*//;')
      if [ "${cluster}" == "" ]
      then
         # just a host to scan and add
         if [ "${fast}" != 'yes' -a "${mode}" != '-' ]
         then
            # test if exists host
            if host ${justhost} | grep -q 'not found'
            then
               [ "${verbose}" == 'yes' ] && echo "Warning: ${justhost} does not exists"
               continue
            fi
            if ! nmap -p 22 -sT -PN ${justhost} | grep -q '\bopen\b'
            then
               if host ${justhost}.${dyn_domain} | grep -q 'not found' || ! nmap -p 22 -sT -PN ${justhost}.${dyn_domain} | grep -q '\bopen\b'
               then
                  [ "${verbose}" == 'yes' ] && echo "Warning: ${justhost} is down"
                  continue
               else
                  [ "${verbose}" == 'yes' ] && echo "Warning: remove ssh key of ${justhost}.${dyn_domain}"
                  host=${justhost}.${dyn_domain}
                  ssh-keygen -q -R $(LANG=C host ${justhost} | awk '{print $4}')
               fi
            fi
         fi
         [ "${verbose}" == 'yes' ] && echo "Warning: add ${host} on list with mode ${mode}"
         echo "${host}" >> "/tmp/${base_path}/master${mode}"
      else
         # cluster, jump in a recursive mode
         [ "${verbose}" == 'yes' ] && echo "Warning: recursive call for cluster ${justhost} (${cluster}), with mode ${mode}"
         cluster=$(get_host_list --mode=${mode} "${cluster}")
      fi
   done
   }
declare -fx get_host_list

get_host_list $@
cat "/tmp/${base_path}/master+" >> "/tmp/${base_path}/master"
for f in $(grep . "/tmp/${base_path}/master-")
do
   egrep "^${f}$" "/tmp/${base_path}/master+" && continue
   echo "${f}" >> "/tmp/${base_path}/master--"
done
for f in $(grep . "/tmp/${base_path}/master")
do
   egrep "^${f}$" "/tmp/${base_path}/master--" && continue
   echo "${f}" >> "/tmp/${base_path}/master++"
done

# split master list in paquet of split_number computer
sort -u "/tmp/${base_path}/master++" | split -l ${split_number} - /tmp/${base_path}/__splitted_

# wait is needed by time tmux session open and ssh time connection
tempo=0.8
first_tempo=0
other_tempo=0
if [ -n "${remote_command}" ]
then
   # add tempo after remote command
   other_tempo=${tempo}
   first_tempo="${tempo} ${other_tempo}"
fi

# loop on each split windows
for f in $(ls -1 /tmp/${base_path}/ | grep ^__splitted_)
do
   if [ "${verbose}" == 'yes' ]
   then
      echo "Info: next hosts to be splitted"
      cat "/tmp/${base_path}/${f}" | sed -e 's/^/  /;'
      sleep 5
   fi

   session=$(shuf -n 1 /usr/share/dict/words | tr -cd "[:alpha:]")

   IFS=$'\n' host=($(cat "/tmp/${base_path}/${f}"))

   tmux -2 new-session -d -s $session "ssh ${host[0]} ${remote_command}; sleep ${first_tempo}"
   # wait ${tempo} second to let new session start...
   sleep ${tempo}

   for (( i=1 ; i < ${#host[@]} ; i++))
   do
      # wait ${tempo} needed in case of ${remote_command}
      tmux splitw -t $session "ssh ${host[$i]} ${remote_command}; sleep ${other_tempo}"
      tmux select-layout tiled
   done

   tmux set-window-option synchronize-panes on  > /dev/null
   tmux set-window-option -g utf8 on            > /dev/null
   tmux set -g default-terminal screen-256color > /dev/null
   #tmux set-option -g set-clipboard on
 
   # Sane scrolling
   #tmux set -g mode-mouse on
   #tmux set -g mouse-resize-pane on
   #tmux set -g mouse-select-pane on
   #tmux set -g mouse-select-window on
 
   #set -g terminal-overrides 'xterm*:smcup@:rmcup@'
 
   # toggle mouse mode to allow mouse copy/paste
   # set mouse on with prefix m
   tmux bind m \
      set -g mode-mouse on \; \
      set -g mouse-select-pane on \; \
      display 'Mouse: ON' > /dev/null
      # set -g mouse-resize-pane on \; \
      #set -g mouse-select-window on \; \
   # set mouse off with prefix M
   tmux bind M \
      set -g mode-mouse off \; \
      set -g mouse-select-pane off \; \
      display 'Mouse: OFF' > /dev/null
      #set -g mouse-resize-pane off \; \
      #set -g mouse-select-window off \; \
   # toggle Broadcast
   tmux bind b set-window-option synchronize-panes

   tmux attach -t $session
done

# Clean temporary folder
[ -d "/tmp/${base_path}" ] && rm -rf "/tmp/${base_path}"