#!/bin/bash
#
# Do basic checks and definitions up front here
#
setupDir=`dirname $0`
if [ -e $setupDir/clusco_setup.sh ] ; then
   source $setupDir/clusco_setup.sh
elif [[ "$clusCoDir" == "" ]] ; then
   echo "Abort: No clusco_setup file found"
   exit 2
fi

#
# Only users who have ACLs for the clusCoDir area can run this
#
touch ${clusCoDir}/${USER}_test.txt > /dev/null 2>&1
if [ $? -ne 0 ] ; then
  echo "*** You are not authorized to run this command ***"
  exit 1
fi
/bin/rm -f ${clusCoDir}/${USER}_test.txt

# Default to no list of machine names
# Also initially set other global variables to empty strings
wn_name_list=""
whichOpt=""
do_this_machine_list=""
pState=""
to_whom=""
# Set up the full list of appliances at each site for the
# T2WN, MSUWN and UMWN wildcards
msuWNappList="cc-*"
umWNappList="bl-*,c4-*,dc-*,dc2-*,sx-*"

# Set this to zero during normal operations
DEBUG=0

#
# Global variables that get defined, and their meanings.
#
# modeInteract     =1 when in fully interactive mode, =0 otherwise
# whichOpt         Command_number that is chosen
# do_this_machine_list  Comma separated list of machines to operate upon
#                  Each entry in the list can have no more than one * wildcard
# do_this_machine  Fully expanded list of machines to work upon
# pState           State ON or OFF for commands requiring this
# to_whom          The Email address for commands requiring it
# privateName      The full private network name of the active machine
# publicName       The full public network name of the active machine
# privateDomain    The private DNS domain of the active machine
# wn_name_list     For commands executing off a list of machines, this is the
#                  name of the file containing the list
#
# ============== Exit codes ==============
# 0 = normal
# 1 = You are root, must run as a normal user
# 2 = Cannot find clusco_setup.sh
# 3 = Unknown command line option
# 4 = Invalid set state, must be ON or OFF
# 5 = Missing required set state
# 6 = Attempt to run push_command in command-line mode
# 7 = Error parsing arguments using getopt
# 8 = Missing required Email address
# 9 = Missing required machine list option
# 10 = Missing required command option

# -----------------------------------------
usage()
{
cat <<EOF
Usage:  cluster_control  Full Interactive invocation, no arguments accepted
        cluster_control [-h|--help] 
                        -c|--command <command number> 
                        -f|--filter machine_names 
                        [-s|--state ON|OFF]
                        [-e|--email email_address] 

          command and filter are always required if any arguments are supplied
          at all, with the exception of the help command

          -h|--help       Print this help text and exit
          -c|--command <command_number>   The command to execute.  run the
                          full interactive command to see the legal list
          -f|--filter machine_names   Execute command_number on these machines.
                          This is a comma separated list, with each entry
                          including up to one wildcard.  There are 3 special
                          wildcards that over-ride any single machine(s);
                          T2WN, MSUWN and UMWN.  The first is exclusive, the
                          last 2 can be combined (but just equal T2WN then)
          -s|--state ON|OFF   The state to set in commands that require
                          this argument.
          -e|--email email_address   The address to which command completion
                          notifications are sent, for commands that require
                          an email address

          Unknown or invalid options cause cluster_control to abort with an
          error code.

EOF
}

# -----------------------------------------
parse_arguments() 
{
TEMP=`getopt -n cluster_control -o c:f:s:e:h \
  -l command:,filter:,state:,email:,help  \
  -- "$@"`
# Die if they fat finger arguments
if [ $? -ne 0 ] ; then
  echo "Error parsing arguments using getopt. Try cluster_control --help"
  exit 7
fi

eval set -- "$TEMP"
while true; do
      case $1 in 
	  -c|--command )
	      whichOpt="$2"; shift; shift; continue
	      ;;
	  -f|--filter )
	      do_this_machine_list="$2"; shift; shift; continue
	      ;;
	  -s|--state )
	      pState="$2"; shift; shift; continue
	      ;;
	  -e|--email )
	      to_whom="$2"; shift; shift; continue
	      ;;
	  -h|--help )
	      usage
	      exit 0
	      ;;
	  --)
	      break
	      ;;
	  *)
	      echo "Unknown option $1\n"
	      exit 3
	      ;;
      esac
done
}

# -----------------------------------------
function_transfer()
{
case $whichOpt in
	1)
	  assumptions ;;
	2)
	  echo "+++ Do a peaceful idle of Condor"
          do_idle ;;
        3)
          echo "+++ Turn on Condor job running"
          do_on ;;
        4)
          echo "+++ Immediately turn off Condor job running"
          do_off ;;
        5)
          echo "+++ Get Power State from DB"
	  get_power_state ;;
        6)
          echo "+++ Set Power State in DB"
	  set_power_state ;;
        7)
          echo "+++ Get Condor State from DB"
          get_condor_state ;;
        8)
          echo "+++ Set Condor State in DB"
          set_condor_state ;;
        9)
	  echo "+++ Use file with list of machines, one per line"
          set_file_name ;;
        10)
          echo "+++ Run an arbitrary command using a machine list"
	  push_command ;;
        11)
          usage ;;
        12)
          echo "+++ Quit"
          cVar=0 ;;
        *)
          echo "        Invalid input" ;;
esac
}

machineInfo ()
{
  cmdAbort=0
  mach=$1
# Need to strip off any domain extensions here for matching purposes
#
  machine=`echo $mach | sed -e s/[.]/\ /g | awk '{print $1}'`

  privateName=$(python $clusCoCmd $machine NAME_PRIVATE)

  if [ "$privateName" == "" ]
      then
      echo "Abort: Named host $mach does not exist."
      cmdAbort=1
  fi

  publicName=$(python $clusCoCmd $machine NAME_PUBLIC)
  privateDomain=$(python $clusCoCmd $machine DOMAIN_PRIVATE)

  publicName=${publicName}.aglt2.org
  privateName=${privateName}.${privateDomain}

  echo $privateName $publicName
  return $cmdAbort
}

#------------------------------------------
make_machine_list()
{
  set -f
  noSquashList=`echo $do_this_machine_list|tr , "\ "`
  if [ $DEBUG -eq 1 ] ; then
    echo "unsquashed list is $noSquashList"
  fi
  do_this_machine=""
  for nMach in $noSquashList ; do
    newTst=`perl $clusCoDir/prlSubStar.pl $nMach`
    if [ $DEBUG -eq 1 ] ; then
      echo " element from unsquashed list is $nMach"
      echo " grep machines using string $newTst"
      echo " Looking in file $clusCoDir/nodeinfo.localnames"
    fi
    while read line; do
      tmpMach=`echo $line | awk -F , '{print $1}' | tr -d \"`
      do_this_machine="${do_this_machine}${tmpMach} "
    done < <(grep -e "$newTst\"" $clusCoDir/nodeinfo.localNames)
  done
  set +f
}

# -----------------------------------------
prompt_machine()
{
  if [ $modeInteract -eq 1 ] ; then
    echo -n "Enter machine name: "
    read do_this_machine_list
  fi
#
# Check for the 3 key names, T2WN, UMWN and MSUWN
#
  echo $do_this_machine_list | grep -q -e T2WN
  if [ $? -eq 0 ] ; then
    do_this_machine_list="$msuWNappList,$umWNappList"
  else
    echo $do_this_machine_list | grep -q -e MSUWN
    if [ $? -eq 0 ] ; then
      hasMSU=1
    else
      hasMSU=0
    fi
    echo $do_this_machine_list | grep -q -e UMWN
    if [ $? -eq 0 ] ; then
      hasUM=2
    else
      hasUM=0
    fi
    (( hasSite = hasUM + hasMSU ))
    case $hasSite in
	1 ) 
	    do_this_machine_list=$msuWNappList
	    ;;
	2 )
	    do_this_machine_list=$umWNappList
	    ;;
	3 )
	    do_this_machine_list="$msuWNappList,$umWNappList"
	    ;;
	* ) ;;
    esac
  fi
    
  if [ $DEBUG -eq 1 ] ; then
    echo "Modifed machine list..."
    echo "  $do_this_machine_list"
  fi
    
  make_machine_list

  if [ $DEBUG -eq 1 ] ; then
    echo "Full list is..."
    echo $do_this_machine
    exit 11
  fi
  return 0
}

# -----------------------------------------
set_file_name()
{
  echo -n "Enter file containing WN list: "
  read newFile
  if [ ! -f $newFile ] ; then
    echo "File does not exist, set no file"
    wn_name_list=""
  else
    wn_name_list="$newFile"
  fi
}

# -----------------------------------------
assumptions()
{
cat <<EOF
=========== Assumptions ============
1. You should have your Kerberos token.
2. If you can see this, you are already an allowed manager.
3. The name of the file that has the list of machines to
   operate upon is remembered until it is either changed, or a
   NULL entry is entered at that prompt.
4. If you are starting Condor running, everything should already have
   been checked OK, eg, lustre mounted, health OK...
5. For commands executed with a list of machines, the file "failed.list"
   in $PWD will detail machines skipped 
   or where the command failed.

EOF
}

# -----------------------------------------
do_idle()
{
  if [ $modeInteract -eq 1 ] ; then
    echo -n "Notification Email: "
    read to_whom
    if [ x$wn_name_list != x ] ; then
      push_cmd.sh -f $wn_name_list -e UP -r "wn_idle $to_whom "
    else
      prompt_machine
      for nMach in $do_this_machine ; do
	  infoList=$(machineInfo "${nMach}")
	  if [ $? -ne 0 ] ; then
	    echo "    Invalid machine name $nMach"
	  fi
	  wn_idle $to_whom $nMach
      done
    fi
  else
    if [ x$to_whom == x ] ; then
      echo "Missing required Email address via -e"
      exit 8
    fi
    prompt_machine
    for nMach in ${do_this_machine} ; do
      wn_idle ${to_whom} ${nMach}
    done
  fi
}

# -----------------------------------------
do_off()
{
  if [ $modeInteract -eq 1 ] ; then
    if [ x$wn_name_list != x ] ; then
      push_cmd.sh -f $wn_name_list -e UP -r "wn_off "
    else
      prompt_machine
      for nMach in $do_this_machine ; do
	  infoList=$(machineInfo "${nMach}")
	  if [ $? -ne 0 ] ; then
	      echo "    Invalid machine name $nMach"
	  fi
	  wn_off ${nMach}
      done
    fi
  else
    prompt_machine
    for nMach in ${do_this_machine} ; do
      wn_off ${nMach}
    done
  fi
}

# -----------------------------------------
do_on()
{
  if [ $modeInteract -eq 1 ] ; then
    if [ x$wn_name_list != x ] ; then
      push_cmd.sh -f $wn_name_list -e UP -r "wn_on "
    else
      prompt_machine
      for nMach in $do_this_machine ; do
	  infoList=$(machineInfo "${nMach}")
	  if [ $? -ne 0 ] ; then
	      echo "    Invalid machine name $nMach"
	  fi
	  wn_on $nMach
      done
    fi
  else
    prompt_machine
    for nMach in ${do_this_machine} ; do
      wn_on ${nMach}
    done
  fi
}

# -----------------------------------------
get_real_power()
{
# Machine name is passed as argument #1
  infoList=$(machineInfo "${1}")
  if [ $? -ne 0 ] ; then
    echo "    Invalid machine name $1"
    return 0
  fi
  power_state=$(get_power_state.sh $1)
  case $power_state in
    0) echo "Power of $1 is OFF" ;;
    1) echo "Power of $1 is ON" ;;
  esac
}

# -----------------------------------------
get_power_state()
{
  if [ x$wn_name_list != x ] ; then
    while read line ; do
	get_real_power ${line}
    done < $wn_name_list
  else
    prompt_machine
    for nMach in $do_this_machine ; do
      get_real_power ${nMach}
    done
  fi
}

# -----------------------------------------
set_real_power()
{
# Provide the machine name, and the power state, as arguments
  infoList=$(machineInfo "${1}")
  if [ $? -ne 0 ] ; then
    echo "    Invalid machine name $1"
    return 0
  fi
  set_power_state.sh $1 $2
}

# -----------------------------------------
set_power_state()
{
  if [ $modeInteract -eq 1 ] ; then
    bVar=1
    while (( bVar == 1 ))
    do
      echo -n "Set power ON or OFF? "
      read pState
      pState=`echo $pState | tr '[a-z]' '[A-Z]'`
      case $pState in
	ON )
           nState="UP"
           bVar=0 ;;
        OFF ) 
	   nState="DOWN"
           bVar=0 ;;
      esac
    done
  else
    if [ x$pState != x ] ; then
      case $pState in
	ON )
           nState="UP"
           ;;
        OFF ) 
	   nState="DOWN"
           ;;
	* )
	   echo "Invalid state $pState"
	   exit 4
	   ;;
      esac
    else
      echo "This command requiires a set state (ON|OFF) via -s"
      exit 5
    fi
  fi
#
  if [ x$wn_name_list != x ] ; then
    while read line ; do
	set_real_power ${line} ${nState}
    done < $wn_name_list
  else
    prompt_machine
    for nMach in $do_this_machine ; do
      set_real_power ${nMach} ${nState}
    done
  fi
}

# -----------------------------------------
get_real_condor_state()
{
  infoList=$(machineInfo "${1}")
  if [ $? -ne 0 ] ; then
    echo "    Invalid machine name $1"
    return 0
  fi
  condor_state=$(get_condor_state.sh $1)
  case $condor_state in
    0) echo "Condor run state of $do_this_machine is OFF" ;;
    1) echo "Condor run state of $do_this_machine is ON" ;;
  esac
}

# -----------------------------------------
get_condor_state()
{
  if [ x$wn_name_list != x ] ; then
    while read line ; do
	get_real_condor_state ${line}
    done < $wn_name_list
  else
    prompt_machine
    for nMach in $do_this_machine ; do
      get_real_condor_state ${nMach}
    done
  fi
}

# -----------------------------------------
set_real_condor_state()
{
  infoList=$(machineInfo "${1}")
  if [ $? -ne 0 ] ; then
    echo "    Invalid machine name $1"
    return 0
  fi
  set_condor_state.sh $1 $2
}

# -----------------------------------------
set_condor_state()
{
  if [ $modeInteract -eq 1 ] ; then
    bVar=1
    while (( bVar == 1 ))
    do
      echo -n "Set condor state in DB ON or OFF? "
      read pState
      pState=`echo $pState | tr '[a-z]' '[A-Z]'`
      case $pState in
	ON )
           nState="UP"
           bVar=0 ;;
        OFF ) 
	   nState="DOWN"
           bVar=0 ;;
      esac
    done
  else
    if [ x$pState != x ] ; then
      case $pState in
	ON )
           nState="UP"
           ;;
        OFF ) 
	   nState="DOWN"
           ;;
	* )
	   echo "Invalid state $pState"
	   exit 4
	   ;;
      esac
    else
      echo "This command requiires a set state (ON|OFF) via -s"
      exit 5
    fi
  fi
#
  if [ x$wn_name_list != x ] ; then
    while read line ; do
	set_real_condor_state ${line} ${nState}
    done < $wn_name_list
  else
    prompt_machine
    for nMach in $do_this_machine ; do
      set_real_condor_state ${nMach} ${nState}
    done
  fi
}

# -----------------------------------------
push_command()
{
  if [ $modeInteract -ne 1 ] ; then
    echo "This command cannot be run (yet) in command-line mode"
    exit 6
  fi
  if [ x$wn_name_list == x ] ; then
    set_file_name
  fi
  if [ x$wn_name_list == x ] ; then
    echo "This command requires a file name for execution"
    return 0
  fi

  isWhere=1
  while (( isWhere == 1 )) 
  do
    echo -n "Execute command remotely (1) or on `hostname -s` (2)? "
    read rem_local
    case $rem_local in
      1) 
        arg="-l"
        isWhere=0 ;;
      2)
        arg="-r"
	isWhere=0 ;;
      *)
    esac
  done

  echo "Choose conditional for the command:"
  echo "   1 Power is ON"
  echo "   2 All machines in the list"
  echo "   Anything else (default) Condor is running"
  echo -n "Choice 1,2 or anything: "
  read macState
  case $macState in
    1)
      conditional="-e UP" ;;
    2)
      conditional="-e NONE" ;;
    *)
      conditional="" ;;
  esac

  echo -n "Command: "
  read pCommand

  push_cmd.sh -f $wn_name_list $conditional $arg "$pCommand"
}

# -----------------------------------------
# 
# OK, here is the MAIN LOOP for this process
# Two methods of invocation, with or without arguments
# If WITHOUT arguments, then loop interactively
# if WITH arguments, these must be provided
# -c command number, 1-11
# -f list of machines.  Single * wildcard is allowed.
#    Special values are WN, UMWN and MSUWN
# -s ON or OFF (some but not all commands require this, no default)
#
if [ $# -ne 0 ] ; then
  modeInteract=0
  parse_arguments $@
#
# These two arguments are required for all commands that have made
# it this far
#
  if [ x$do_this_machine_list == x ] ; then
    echo "Missing required machine list via -f"
    exit 9
  fi
  if [ x$whichOpt == x ] ; then
    echo "Missing required command via -c"
    exit 10
  fi
#
  function_transfer
  exit 0
else
#
# Interactive invocation
#
  modeInteract=1
  cVar=1
  while (( cVar == 1 )) ; do

    cat <<EOF

*****************
Choose an Option
   1. What are the assumptions here?
   2. Peacefully Idle a Worker Node
   3. Start up Condor on a Worker Node
   4. Immediately Stop Condor on a Worker Node
   5. Get Power state of a Worker Node from the DB
   6. Change DB Power state of a Worker Node
   7. Get Condor state of a Worker Node from the DB
   8. Change DB Condor state of a Worker Node
   9. If doing many machines, name of file containing list [$wn_name_list]
   10. Execute a command on the machine list above
   11. Print help text
   12. Quit
EOF
    echo -n "Option: "
    read whichOpt

    function_transfer

  done

fi