#!/bin/bash
#
# 2014/03/26 Gabriel Moreau
#
# 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}"' QUIT

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

AUTHOR
   Gabriel Moreau

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

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: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
         ;;
      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
   for host in $*
   do
      mode=""
      last_char="${host: -1}"
      if [ "${last_char}" == "-" -o "${last_char}" == "+" ]
      then
         mode="${last_char}"
         host="${host:0:${#host}-1}"
      fi

      cluster=$(grep "^${host}" ${HOME}/.csshrc | cut -f 2 -d '=')
      if [ "${cluster}" == "" ]
      then
         if [ "${fast}" != 'yes' ]
         then
            # short host without login part if any
            if host ${host#*@} | grep -q 'not found'
            then
               [ "${verbose}" == "yes" ] && echo Warning: ${host#*@} does not exists
               continue
            fi
            if ! nmap -p 22 -sT -PN ${host#*@} | grep -q '\bopen\b'
            then
               if host ${host#*@}.${dyn_domain} | grep -q 'not found' || ! nmap -p 22 -sT -PN ${host#*@}.${dyn_domain} | grep -q '\bopen\b'
               then
                  [ "${verbose}" == "yes" ] && echo Warning: ${host#*@} is down
                  continue
               else
                  host=${host}.${dyn_domain}
                  ssh-keygen -q -R $(LANG=C host ${host} | awk '{print $4}')
               fi
            fi
         fi
         echo "${host}" >> "/tmp/${base_path}/master${mode}"
      else
         cluster=$(get_host_list "${cluster}${mode}")
      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

sort -u "/tmp/${base_path}/master++" | split -l ${split_number} - /tmp/${base_path}/_

for f in $(ls -1 /tmp/${base_path}/ | grep -v ^master)
do
   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]}"


   for (( i=1 ; i < ${#host[@]} ; i++))
   do
      tmux splitw -t $session "ssh ${host[$i]}"
      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}"
