#!/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] <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
   -h  help

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...

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}"