#!/bin/bash
# Name: send-cmd.sh
# Description: Execute a command(s) on the specified system(s)
# Last Updated: 2016-12-20
# Recent Changes:-Added interactive mode (-i)
#                -Added worker argument and return code for ssh commands; newline
#                 after output; color to return code
####################################################################################

##### Customize These Variables #####
# Spacecmd command and any default arguments
spacecmd_cmd="spacecmd -q"

# Max number of workers at a time
max_workers=10

# Get base path
base_path="$(echo ${0} | sed 's/send-cmd.sh//')"

# Worker script
worker_script="${base_path}worker_send-cmd.sh"

## Define colors ##
# End/reset color
color_end='\033[0m'

# Colors
color_green='\033[0;32m'
color_red='\033[0;31m'
color_yellow='\033[0;33m'
##### End of Customize Variables #####

#=====================================
# Functions; Main starts after
#=====================================
function show_usage
{
  echo -e "\n==== Send Command ===="
  echo -e "\nDescription: Execute a command on the specified system(s)."
  echo -e "\n--Usage--"
  echo -e "send-cmd.sh <SYSTEM-OPTION> -c 'command to send' [OTHER OPTIONS]"
  echo -e "\nSYSTEM OPTIONS"
  echo -e "-s system             => Single system hostname(unquoted)."
  echo -e "-s 'system1 system2'  => Multiple system hostnames(quoted)."
  echo -e "-s filename           => File with system name(s)."
  echo -e "-g system-group       => Spacewalk system group."
  echo -e "-a                    => All Spacewalk registered systems."
  echo -e "\nCOMMAND SYNTAX"
  echo -e "-c 'command to send'  => Send the quoted command.(Can be unquoted if only 1 word)"
  echo -e "\nOTHER OPTIONS"
  echo -e "-h                    => Display usage."
  echo -e "-i                    => Interactive mode.(Default: Non-interactive)"
  echo -e "-p                    => Send commands in parallel."
  echo -e "-w num                => Set the max workers for parallel mode (default:${max_workers})"
  echo -e "-v                    => Verbose output."
  echo -e "\n--Other Requirements--"
  echo -e "-> Spacecmd and config file setup required.(-g or -a only)"
  echo -e "-> Run as a user with system list privileges on the Spacewalk server.(-g or -a only)"
  echo -e
}

#=======================
# Get Script Arguments
#=======================
# Reset POSIX variable in case it has been used previously in this shell
OPTIND=1

## Default settings ##
# Non-interactive by default
interactive_mode="no"

# Do not send to all systems by default
all_systems="no"

# Send commands in serial by default
parallel_cmds="no"

# Verbose to no by default
verbose="no"

## Get command line arguments ##
while getopts "his:g:ac:pw:v" opt; do
  case "${opt}" in
    h) # -h (help) argument
      show_usage
      exit 0
    ;;
    i) # -i (interactive) argument
      interactive_mode="yes"
    ;;
    s) # -s system
      system_name="${OPTARG}"
      # Determine if cmd type is single system(s) argument or filename
      if [[ -f ${system_name} ]]; then
        cmd_type="file"
      else
        cmd_type="single"
      fi
    ;;
    g) # -g system-group
      system_group="${OPTARG}"
      cmd_type="group"
    ;;
    a) # -a (all systems)
      all_systems="yes"
      cmd_type="all"
    ;;
    c) # command to send
      send_cmd="${OPTARG}"
    ;;
    p) # send commands in parallel
      parallel_cmds="yes"
    ;;
    w) # max workers in parallel
      max_workers="${OPTARG}"
    ;;
    v) # verbose output
      verbose="yes"
    ;;
    *) # invalid argument
      show_usage
      exit 0
    ;;
  esac
done

## Argument Sanity Checks ##
# Scenario: No arguments set
if [[ ${all_systems} == "no" ]]; then
  if [[ -z ${system_name} && -z ${system_group} ]]; then
    echo -e "\n>> ERROR! You must decide on what system(s) will be sent the command."
    show_usage
    exit 1
  fi
fi

# Scenario: All systems AND single system(s) argument/filename set
if [[ ${all_systems} == "yes" && ${system_name} ]]; then
  echo -e "\n>> ERROR! Incompatible arguments (all systems and a single system(s) argument/file)."
  show_usage
  exit 1
fi

# Scenario: All systems AND system group set
if [[ ${all_systems} == "yes" && ${system_group} ]]; then
  echo -e "\n>> ERROR! Incompatible arguments (all systems and a system group)."
  show_usage
  exit 1
fi

# Scenario: Single system(s) argument AND system group set
if [[ ${system_name} && ${system_group} ]]; then
  echo -e "\n>> ERROR! Incompatible arguments (single system(s) argument/file and a system group)."
  show_usage
  exit 1
fi

# Scenario: No command set
if [[ -z ${send_cmd} ]]; then
  echo -e "\n>> ERROR! You must enter a command to send."
  show_usage
  exit 1
fi

# Parallel Send and Interactive Mode set: Issue warning
if [[ ${parallel_cmds} == "yes" && ${interactive_mode} == "yes" ]]; then
  echo -e "\n>> WARNING: Interactive mode (-i) ignored due to parallel send mode (-p).\n"
  interactive_mode="no"
fi

## Command Type: Single or Filename ##

#===================
# Pre-checks: Ensure dependencies exist
#===================

# Only check for Spacewalk dependencies if NOT sending to a single system(s) argument/file
if [[ -z ${system_name} ]]; then
  # Check for Spacecmd
  which spacecmd &> /dev/null
  if [[ $? -ne 0 ]]; then
    echo "\n>> Error! The command 'spacecmd' is not found or not in PATH. Exiting..."
    exit 1
  fi

  # Check to see if a spacecmd config file exists
  if [[ ! -f ${HOME}/.spacecmd/config ]]; then
    echo -e "\n>> Error! No spacecmd config file found at: ${HOME}/.spacecmd/config. Exiting..."
    exit 1
  fi
fi

#===================
# Main starts here
#===================

if [[ ${verbose} == "yes" ]]; then
  echo -e "============================="
  echo -e "####=== Send Command ====####"
  echo -e "============================="

  echo -e "NOTE: Commands with spaces and multiple system names must be quoted.\n"

  if [[ ${system_name} ]]; then
    if [[ ${cmd_type} == "file" ]]; then
      echo -e "Send command to these system(s) from file(${system_name}): \n$(cat ${system_name})"
    else
      echo -e "Send command to system(s): ${system_name}"
    fi
  elif [[ ${system_group} ]]; then
    echo -e "Send command to systems in this Spacewalk group: ${system_group}"
  else
    echo -e "Send command to ALL systems."
  fi

  if [[ ${parallel_cmds} == "yes" ]]; then
    echo -e "Send Mode: Parallel with (${max_workers}) workers"
  else
    echo -e "Send Mode: Serial"
  fi

  echo -e "Command to send: ${send_cmd}"
  echo -e "=>Continue?[y/n]:\c"
  read run_script

  if [[ ${run_script} != "y" ]]; then
    echo -e "\n>>Will not run the send command script. Exiting..."
    exit 1
  fi
fi # end of verbose check

# If we are using parallel commands, set the current number of workers
if [[ ${parallel_cmds} == "yes" ]]; then
  current_workers=0
fi

case ${cmd_type} in
  single)
  ## Single system(s) argument ##
  if [[ ${verbose} == "yes" ]]; then
    echo -e "\n>> Sending command(s) to system(s)..."
  fi

  if [[ ${parallel_cmds} == "yes" ]]; then
    # Parallel Execution
    for node in $(echo ${system_name}); do

      # If the current number of workers equals the max, wait for them to complete, then reset to zero
      if [[ ${current_workers} -ge ${max_workers} ]]; then
        wait
        current_workers=0
      fi

      # Start a new worker in the background
      if [[ ${verbose} == "yes" ]]; then
        echo "-> Working on ${node}..."
      fi
      (${worker_script} ${node} "${send_cmd}") &

      # Increase the number of current workers
      current_workers=$(( ${current_workers} + 1 ))
    done
    # Wait for all remaining workers to complete
    wait
  else
    # Serial Execution
    for node in $(echo ${system_name}); do

      # Non-interactive (default)
      if [[ ${interactive_mode} != "yes" ]]; then
        output="$(ssh -qt -o ConnectTimeout=5 ${node} "${send_cmd}")"
        return_code=$(echo $?)
        case "${return_code}" in
          0) # 0 return code - show green return code
            echo -e "-> ${node} (${color_green}retcode=${return_code}${color_end})\n${output}\n"
          ;;
          1) # 1 return code  - show red return code
            echo -e "-> ${node} (${color_red}retcode=${return_code}${color_end})\n${output}\n"
          ;;
          *) # any other return code  - show yellow return code
            echo -e "-> ${node} (${color_yellow}retcode=${return_code}${color_end})\n${output}\n"
          ;;
        esac

      else
        # Interactive mode set
        echo -e "-> ${node}"
        ssh -qt -o ConnectTimeout=5 ${node} "${send_cmd}"
        echo
      fi
    done
  fi

  ;;
  file)
  ## File with one or more systems ##
  if [[ ${verbose} == "yes" ]]; then
    echo -e "\n>> Sending command(s) to system(s) in file (${system_name})..."
  fi

  if [[ ${parallel_cmds} == "yes" ]]; then
    # Parallel Execution
    for node in $(cat ${system_name}); do

      # If the current number of workers equals the max, wait for them to complete, then reset to zero
      if [[ ${current_workers} -ge ${max_workers} ]]; then
        wait
        current_workers=0
      fi

      # Start a new worker in the background
      if [[ ${verbose} == "yes" ]]; then
        echo "-> Working on ${node}..."
      fi
      (${worker_script} ${node} "${send_cmd}") &

      # Increase the number of current workers
      current_workers=$(( ${current_workers} + 1 ))
    done
    # Wait for all remaining workers to complete
    wait
  else
    # Serial Execution
    for node in $(cat ${system_name}); do

      # Non-interactive (default)
      if [[ ${interactive_mode} != "yes" ]]; then
        output="$(ssh -qt -o ConnectTimeout=5 ${node} "${send_cmd}")"
        return_code=$(echo $?)
        case "${return_code}" in
          0) # 0 return code - show green return code
            echo -e "-> ${node} (${color_green}retcode=${return_code}${color_end})\n${output}\n"
          ;;
          1) # 1 return code  - show red return code
            echo -e "-> ${node} (${color_red}retcode=${return_code}${color_end})\n${output}\n"
          ;;
          *) # any other return code  - show yellow return code
            echo -e "-> ${node} (${color_yellow}retcode=${return_code}${color_end})\n${output}\n"
          ;;
        esac

      else
        # Interactive mode set
        echo -e "-> ${node}"
        ssh -qt -o ConnectTimeout=5 ${node} "${send_cmd}"
        echo
      fi
    done
  fi

  ;;
  group)
  ## Group of systems (Spacewalk Group) ##
  if [[ ${verbose} == "yes" ]]; then
    echo -e "\n>> Sending command(s) to a group of systems (${system_group})..."
  fi

  # Check to see if the Spacewalk group exists; exit if it does not
  ${spacecmd_cmd} group_list | grep ${system_group} > /dev/null
  if [[ $? -ne 0 ]]; then
    echo -e "-> ERROR! Could not find Spacewalk group: ${system_group}"
    exit 1
  fi

  if [[ ${parallel_cmds} == "yes" ]]; then
    # Parallel Execution
    for node in $(${spacecmd_cmd} group_listsystems ${system_group}); do

      # If the current number of workers equals the max, wait for them to complete, then reset to zero
      if [[ ${current_workers} -ge ${max_workers} ]]; then
        wait
        current_workers=0
      fi

      # Start a new worker in the background
      if [[ ${verbose} == "yes" ]]; then
        echo "-> Working on ${node}..."
      fi
      (${worker_script} ${node} "${send_cmd}") &

      # Increase the number of current workers
      current_workers=$(( ${current_workers} + 1 ))
    done
    # Wait for all remaining workers to complete
    wait
  else
    # Serial Execution
    for node in $(${spacecmd_cmd} group_listsystems ${system_group}); do

      # Non-interactive (default)
      if [[ ${interactive_mode} != "yes" ]]; then
        output="$(ssh -qt -o ConnectTimeout=5 ${node} "${send_cmd}")"
        return_code=$(echo $?)
        case "${return_code}" in
          0) # 0 return code - show green return code
            echo -e "-> ${node} (${color_green}retcode=${return_code}${color_end})\n${output}\n"
          ;;
          1) # 1 return code  - show red return code
            echo -e "-> ${node} (${color_red}retcode=${return_code}${color_end})\n${output}\n"
          ;;
          *) # any other return code  - show yellow return code
            echo -e "-> ${node} (${color_yellow}retcode=${return_code}${color_end})\n${output}\n"
          ;;
        esac

      else
        # Interactive mode set
        echo -e "-> ${node}"
        ssh -qt -o ConnectTimeout=5 ${node} "${send_cmd}"
        echo
      fi
    done
  fi

  ;;
  all)
  ## All Systems ##
  if [[ ${verbose} == "yes" ]]; then
    echo -e "\n>> Sending command(s) to All systems..."
  fi

  if [[ ${parallel_cmds} == "yes" ]]; then
    # Parallel Execution
    for node in $(${spacecmd_cmd} system_list); do 

      # If the current number of workers equals the max, wait for them to complete, then reset to zero
      if [[ ${current_workers} -ge ${max_workers} ]]; then
        wait
        current_workers=0
      fi

      # Start a new worker in the background
      if [[ ${verbose} == "yes" ]]; then
        echo "-> Working on ${node}..."
      fi
      (${worker_script} ${node} "${send_cmd}") &

      # Increase the number of current workers by one
      current_workers=$(( ${current_workers} + 1 ))
    done
    # Wait for all remaining workers to complete
    wait
  else
    # Serial Execution
    for node in $(${spacecmd_cmd} system_list); do 
      # Non-interactive (default)
      if [[ ${interactive_mode} != "yes" ]]; then
        output="$(ssh -qt -o ConnectTimeout=5 ${node} "${send_cmd}")"
        return_code=$(echo $?)
        case "${return_code}" in
          0) # 0 return code - show green return code
            echo -e "-> ${node} (${color_green}retcode=${return_code}${color_end})\n${output}\n"
          ;;
          1) # 1 return code  - show red return code
            echo -e "-> ${node} (${color_red}retcode=${return_code}${color_end})\n${output}\n"
          ;;
          *) # any other return code  - show yellow return code
            echo -e "-> ${node} (${color_yellow}retcode=${return_code}${color_end})\n${output}\n"
          ;;
        esac

      else
        # Interactive mode set
        echo -e "-> ${node}"
        ssh -qt -o ConnectTimeout=5 ${node} "${send_cmd}"
        echo
      fi
    done
  fi

  ;;
esac

if [[ ${verbose} == "yes" ]]; then
  echo -e "\n============================="
  echo -e "=- Send Command Completed. -="
  echo -e "============================="
fi