#!/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 2;' QUIT INT TERM

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

VERSION='0.3.1'

function usage() {
   cat <<END_USAGE
tssh - tmux cluster ssh

 tssh [-w number] [-f] [-v] [-h] [-p] [-d delay_time] [-c remote_cmd] [-o ssh_option] <host1> <host2> <clusterssh class>... <hostM>- <hostN>+

   -w windows     windows to open (integer, default 16)
   -o ssh_option  option to pass to ssh
   -f             fast, no nmap scan to eliminate sleeping computer
   -v             verbose
   -d delay_time  delay between parallel connexion to remote host
   -c remote_cmd  launch the remote command on hosts and exit
   -p             launch in parallel (only with -c) with command parallel or xargs
   -h             help

VERSION ${VERSION}
END_USAGE
   }

export delay_time=0.8
export remote_command=''
export ssh_option=''
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}/.config/tssh/config.sh" ]
then
   . "${HOME}/.config/tssh/config.sh"
elif [ -e "${HOME}/.tsshrc" ]
then
   . "${HOME}/.tsshrc"
fi

# get options
if [ $# -eq 0 ]; then usage; exit 1; fi 
while getopts "w:o:d:c:fvph" options
do
   case ${options} in
      w)
         if echo ${OPTARG} | egrep -q '^[[:digit:]]+$' && [ ${OPTARG} -gt 0 ]
         then
            export split_number=${OPTARG}
         else
            usage
            exit 4
         fi
         ;;
      d)
         if echo ${OPTARG} | egrep -q '^[[:digit:]]+$' && [ ${OPTARG} -gt 0 ]
         then
            export delay_time=${OPTARG}
         else
            usage
            exit 5
         fi         
         ;;
      c)
         if echo ${OPTARG} | egrep -q '[[:alpha:]]'
         then
            export remote_command=${OPTARG}
         else
            usage
            exit 6
         fi         
         ;;
      p)
         export parallel='yes'
         ;;
      o)
         if echo ${OPTARG} | egrep -q '[[:alpha:][:digit:]]'
         then
            export ssh_option=${OPTARG}
         else
            usage
            exit 7
         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

if [ "${parallel}" == 'yes' -a -n "${remote_command}" ]
then
   if parallel --version 2> /dev/null | grep -q ^'GNU parallel'
   then
      sort -u "/tmp/${base_path}/master++" | parallel -j ${split_number} -I {} -- "ssh ${ssh_option} {} ${remote_command}"
   else
      sort -u "/tmp/${base_path}/master++" | xargs -r -n 1 -P ${split_number} -I {} -- ssh ${ssh_option} {} "${remote_command}"
   fi

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


# 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
first_delay_time=0
other_delay_time=0
if [ -n "${remote_command}" ]
then
   # add delay time after remote command
   other_delay_time=${delay_time}
   first_delay_time="${delay_time} ${other_delay_time}"
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 ${ssh_option} ${host[0]} ${remote_command}; sleep ${first_delay_time}"
   # wait ${delay_time} second to let new session start...
   sleep ${delay_time}

   for (( i=1 ; i < ${#host[@]} ; i++))
   do
      # wait ${delay_time} needed in case of ${remote_command}
      tmux splitw -t $session "ssh ${ssh_option} ${host[$i]} ${remote_command}; sleep ${other_delay_time}"
      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}"


exit


################################################################
# Documentation in POD format (like Perl)
################################################################

=head1 NAME

tssh - tmux cluster ssh

=head1 SYNOPSIS

 tssh [-w number] [-f] [-v] [-h] [-p] [-d delay_time] [-c remote_cmd] [-o ssh_option] <host1> <host2> <clusterssh class>... <hostM>- <hostN>+

=head1 OPTIONS

=head2 C<-w number> 

Windows to open (integer, default 16)

=head2 C<-o ssh_option>

Option to pass to ssh

=head2 C<-f>

Fast, no nmap scan to eliminate sleeping computer

=head2 C<-d delay_time>

When use C<tmux> context,
add a small delay in second before lauching the second shell
to let C<tmux> the time to open the window (default 0.8).

=head2 C<-c remote_cmd>

Launch the remote command on hosts and exit

=head2 C<-p>

Launch in parallel (only with option C<-c>) with command parallel or xargs

=head2 C<-v>

Verbose

=head2 C<-h>

Minimal help



=head1 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 C<host>, C<login@host>, C<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 C<tmux> is C<Ctrl^b>.
You can switch from broadcast to a local machine with C<Ctrl^b Ctrl^b>
and move between machine with C<Ctrl^b ArrowKey>.

=head1 DEPENDS

On Debian, you need the package

 apt-get install tmux ncurses-bin wamerican nmap

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

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

By default, C<tssh> use C<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 C<tput> is not installed, the default is 16 windows...

=head1 CONFIGURATION

The clusterssh config file F<~/.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 C<tssh> config file (F<~/.config/tssh/config.sh>) can be use change the default parameters.

 #export delay_time=0.8
 #export split_number=16
 export dyn_domain='mycompagny.local'
 #export ssh_option=''
 #export fast='yes'
 #export verbose='yes'

=head2 C<delay_time>

See option C<-d>.

=head2 C<split_number>

See option C<-w>.

=head2 C<dyn_domain>

If computer doesn't respond on it's normal IP,
search also to connect on it with this DNS domain.
Could be use with DNS dynamic domain where computer register inside automatically.

=head2 C<ssh_option>

See option C<-o>.

=head2 C<fast>

See option C<-f>.

=head2 C<verbose>

See option C<-v>.



=head1 SEE ALSO

cluster-ssh, cssh, xargs, parallel, tmux

Own project ressources:

=over

=item * L<Web site|http://servforge.legi.grenoble-inp.fr/projects/soft-trokata/wiki/SoftWare/Tssh>

=item * L<Online Manual|http://servforge.legi.grenoble-inp.fr/pub/soft-trokata/tssh/tssh.html>

=item * L<SVN repository|http://servforge.legi.grenoble-inp.fr/svn/soft-trokata/trunk/tssh>

=back


=head1 AUTHORS

Written by Gabriel Moreau, Grenoble - France


=head1 COPYRIGHT

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