peg

#!/bin/bash

# Copyright 2015 Insight Data Science
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

set -o errexit

PEG_ROOT=$(dirname "${BASH_SOURCE}")
AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:=us-west-2}
REM_USER=${REM_USER:=ubuntu}

source ${PEG_ROOT}/util.sh

nargs="$#"

if [ ${nargs} == 0 ]; then
  echo "Usage: peg [options] <command> [parameters]"
  echo "peg: error: too few arguments"
  echo "use the -h option for available commands"
  exit 1
fi

while getopts ":h" opt; do
  case ${opt} in
    h)
      echo "peg"
      echo "    - config"
      echo "    - aws"
      echo "    - validate"
      echo "    - fetch"
      echo "    - describe"
      echo "    - up"
      echo "    - down"
      echo "    - install"
      echo "    - uninstall"
      echo "    - service"
      echo "    - ssh"
      echo "    - sshcmd-node"
      echo "    - sshcmd-cluster"
      echo "    - scp"
      echo "    - retag"
      echo "    - start"
      echo "    - stop"
      echo "    - port-forward"
      exit 1
      ;;

    *)
      echo "Invalid option flag: -${OPTARG}"
      exit 1
      ;;
  esac
done

args=("$@")
command=${args[0]}
parameters=${args[@]:1}
nfargs=$(echo ${parameters} | wc -w)

case ${command} in
  config)
    if [ -z ${AWS_ACCESS_KEY_ID} ]; then
      echo -e "${color_red}[MISSING] AWS_ACCESS_KEY_ID${color_norm}"
    else
      echo "access key: ${AWS_ACCESS_KEY_ID}"
    fi

    if [ -z ${AWS_SECRET_ACCESS_KEY} ]; then
      echo -e "${color_red}[MISSING] AWS_SECRET_ACCESS_KEY${color_norm}"
    else
      echo "secret key: ${AWS_SECRET_ACCESS_KEY}"
    fi

    if [ -z ${AWS_DEFAULT_REGION} ]; then
      echo -e "${color_red}MISSING AWS_DEFAULT_REGION${color_norm}"
    else
      echo "    region: ${AWS_DEFAULT_REGION}"
    fi

    if [ -z ${REM_USER} ]; then
      echo -e "${color_red}MISSING REM_USER${color_norm}"
    else
      echo "  SSH User: ${REM_USER}"
    fi
    ;;

  aws)
    PARAMETER_ARR=(${parameters[@]})
    PEG_AWS_CMD=${PARAMETER_ARR[0]}
    VPC_NAME=${PARAMETER_ARR[1]}
    case ${PEG_AWS_CMD} in
      vpcs)
        show_all_vpcs
        ;;

      subnets)
        show_all_subnets ${VPC_NAME}
        ;;

      security-groups)
        show_all_security_groups ${VPC_NAME}
        ;;

      *)
        echo "specify command for peg aws"
        echo "peg aws"
        echo "        - vpcs"
        echo "        - subnets"
        echo "        - security-groups"
        exit 1
    esac
    ;;

  validate)
    if [[ "${nfargs}" -eq "1" ]]; then
      YAML_FILE=${parameters}
      if [ ! -f ${YAML_FILE} ]; then
        echo "${YAML_FILE} not found"
        exit 1
      else
        set_launch_config ${YAML_FILE}
        validate_template
        echo "Template validated successfully."
      fi
    else
      echo "Invalid number of arguments"
      exit 1
    fi
    ;;

  fetch)
    if [[ "${nfargs}" -eq "1" ]]; then
      CLUSTER_NAME=${parameters}

      CLUSTER_INFO_PATH=${PEG_ROOT}/tmp/${CLUSTER_NAME}
      if [ -d ${CLUSTER_INFO_PATH} ]; then
        rm -rf ${CLUSTER_INFO_PATH}
      fi

      mkdir -p ${CLUSTER_INFO_PATH}

      store_public_dns ${CLUSTER_NAME}
      store_hostnames ${CLUSTER_NAME}
      store_roles ${CLUSTER_NAME}
      store_pemkey ${CLUSTER_NAME}
      peg describe ${CLUSTER_NAME}
    else
      echo "Invalid number of arguments"
      echo "Usage: peg fetch <cluster-name>"
      exit 1
    fi
    ;;

  describe)
    if [[ "${nfargs}" -eq "1" ]]; then
      CLUSTER_NAME=${parameters}

      describe_cluster ${CLUSTER_NAME}
    else
      echo "Invalid number of arguments"
      echo "Usage: peg describe <cluster-name>"
      exit 1
    fi
    ;;

  up)
    if [[ "${nfargs}" -eq "1" ]]; then
      TEMPLATE_PATH=${parameters}

      peg validate ${TEMPLATE_PATH}
      set_launch_config ${TEMPLATE_PATH}
      run_instances
    else
      echo "Invalid number of arguments"
      echo "Usage: peg up <template-path>"
      exit 1
    fi
    ;;

  install)
    if [[ "${nfargs}" -eq "2" ]]; then
      PARAMETER_ARR=(${parameters})
      CLUSTER_NAME=${PARAMETER_ARR[0]}
      TECHNOLOGY=${PARAMETER_ARR[1]}

      check_cluster_exists ${CLUSTER_NAME}

      MASTER_DNS=$(fetch_public_dns_of_node_in_cluster ${CLUSTER_NAME} 1)
      DEP_ROOT_FOLDER=/usr/local/

      case ${TECHNOLOGY} in
        aws)
          echo "Passing AWS credentials to ${CLUSTER_NAME}"
          ${PEG_ROOT}/config/pass_aws_cred ${CLUSTER_NAME}
          ;;

        environment)
          echo "Setting up base environment packages on ${CLUSTER_NAME}"
          ${PEG_ROOT}/install/environment/install_env_cluster.sh ${CLUSTER_NAME}
          ;;

        alluxio|cassandra|elasticsearch|flink|hadoop|hbase|hive|kafka|opscenter|pig|presto|redis|spark|storm|zookeeper)
          get_dependencies >> /dev/null
          install_tech "cluster"
          ;;

        kafka-manager)
          ${PEG_ROOT}/install/kafka-manager/install_kafka_manager_cluster.sh ${CLUSTER_NAME}
          ${PEG_ROOT}/config/kafka-manager/setup_cluster.sh ${CLUSTER_NAME}
          ;;

        kibana)
          get_dependencies >> /dev/null
          install_tech "single"
          ;;

        memsql)
          ${PEG_ROOT}/install/memsql/install_memsql_cluster.sh ${CLUSTER_NAME}
          ;;

        riak)
          ${PEG_ROOT}/install/riak/install_riak_cluster.sh ${CLUSTER_NAME}
          ${PEG_ROOT}/config/riak/setup_cluster.sh ${CLUSTER_NAME}
          ;;

        secor)
          ${PEG_ROOT}/install/secor/install_secor_cluster.sh ${CLUSTER_NAME}
          ${PEG_ROOT}/config/secor/setup_cluster.sh ${CLUSTER_NAME}
          ;;

        ssh)
          echo "Setting up passwordless SSH on ${CLUSTER_NAME}"
          ${PEG_ROOT}/config/ssh/setup_passwordless_ssh.sh ${CLUSTER_NAME}
          ;;

        zeppelin)
          echo "Installing Zeppelin on Master Node in ${CLUSTER_NAME}"
          ${PEG_ROOT}/install/zeppelin/install_zeppelin_cluster.sh ${CLUSTER_NAME}
          ${PEG_ROOT}/config/zeppelin/setup_cluster.sh ${CLUSTER_NAME}
          ;;

        *)
          echo "Invalid technology to install."
          exit 1
          ;;

      esac
    else
      echo "Invalid number of arguments"
      echo "Usage: peg install <cluster-name> <technology>"
    fi
    ;;

  uninstall)
    if [[ "${nfargs}" -eq "2" ]]; then
      PARAMETER_ARR=(${parameters})
      CLUSTER_NAME=${PARAMETER_ARR[0]}
      TECHNOLOGY=${PARAMETER_ARR[1]}

      check_cluster_exists ${CLUSTER_NAME}

      PUBLIC_DNS=($(fetch_cluster_public_dns ${CLUSTER_NAME}))
      MASTER_DNS=${PUBLIC_DNS[0]}

      PEMLOC=${PEG_ROOT}/tmp/$CLUSTER_NAME/*.pem
      ROOT_FOLDER=/usr/local/

      case ${TECHNOLOGY} in
        alluxio|cassandra|elasticsearch|flink|hadoop|hbase|hive|kafka|opscenter|pig|presto|redis|spark|storm|zeppelin|zookeeper)
          INSTALLED=$(check_remote_folder ${MASTER_DNS} ${ROOT_FOLDER}${TECHNOLOGY})
          if [ "$INSTALLED" = "installed" ]; then
            peg service ${CLUSTER_NAME} ${TECHNOLOGY} stop
            uninstall_tech "cluster"
            echo -e "${color_green}${TECHNOLOGY} uninstalled!${color_norm}"
          else
            echo "${TECHNOLOGY} is not installed in ${ROOT_FOLDER}"
            exit 1
          fi
          ;;

        secor|kafka-manager|kibana)
          INSTALLED=$(check_remote_folder ${MASTER_DNS} ${ROOT_FOLDER}${TECHNOLOGY})
          if [ "$INSTALLED" = "installed" ]; then
            peg service ${CLUSTER_NAME} ${TECHNOLOGY} stop
            uninstall_tech "single"
            echo "${TECHNOLOGY} uninstalled!"
          else
            echo "${TECHNOLOGY} is not installed in ${ROOT_FOLDER}"
            exit 1
          fi
          ;;

        memsql)
          INSTALLED=$(check_remote_folder ${MASTER_DNS} ${ROOT_FOLDER}${TECHNOLOGY})
          if [ "$INSTALLED" = "installed" ]; then
            peg service ${CLUSTER_NAME} ${TECHNOLOGY} start
            peg sshcmd-node ${CLUSTER_NAME} 1 "
              . ~/.profile;
              yes DELETE | sudo memsql-ops memsql-delete --all;
              yes DELETE | sudo memsql-ops agent-uninstall --all;
              sudo rmdir /usr/local/memsql;
              sudo rmdir /usr/local/memsql-ops-data;
              sudo rmdir /var/lib/memsql-ops;
              sudo rm -r /home/ubuntu/memsql-ops*;
            "
          fi
          ;;

        riak)
          INSTALLED=$(check_remote_file ${MASTER_DNS} "/usr/sbin/riak")
          if [ "$INSTALLED" = "installed" ]; then
            for dns in ${PUBLIC_DNS[@]}; do
              run_cmd_on_node ${dns} sudo dpkg --purge riak
              run_cmd_on_node ${dns} sudo rm -rf /var/lib/riak
            done
          fi
          ;;

        *)
          echo "Invalid technology to uninstall."
          exit 1
          ;;

      esac
    else
      echo "Invalid number of arguments"
      echo "Usage: peg uninstall <cluster-name> <technology>"
    fi
    ;;

  service)
    if [[ "${nfargs}" -eq "3" ]]; then
      PARAMETER_ARR=(${parameters})
      CLUSTER_NAME=${PARAMETER_ARR[0]}
      TECHNOLOGY=${PARAMETER_ARR[1]}
      ACTION=${PARAMETER_ARR[2]}

      check_cluster_exists ${CLUSTER_NAME}

      MASTER_DNS=$(fetch_public_dns_of_node_in_cluster ${CLUSTER_NAME} 1)

      case $TECHNOLOGY in
        alluxio|cassandra|elasticsearch|flink|hadoop|hbase|kafka|kafka-manager|kibana|memsql|opscenter|presto|redis|secor|spark|storm|zeppelin|zookeeper)
          ROOT_FOLDER=/usr/local/
          INSTALL_PATH="folder"
          service_action
          ;;

        riak)
          ROOT_FOLDER=/usr/sbin/
          INSTALL_PATH="file"
          service_action
          ;;

        *)
          echo "Invalid service to ${ACTION}."
          exit 1
          ;;

      esac
    else
      echo "Invalid number of arguments"
      echo "Usage: peg service <cluster-name> <technology> <start|stop>"
    fi
    ;;

  down)
    if [[ "${nfargs}" -eq "1" ]]; then
      CLUSTER_NAME=${parameters}

      check_cluster_exists ${CLUSTER_NAME}

      INSTANCE_IDS=($(get_instance_ids_with_name_and_role ${CLUSTER_NAME}))
      echo -e "Are you sure you want to terminate ${color_blue}$CLUSTER_NAME${color_norm}?"
      for ids in ${INSTANCE_IDS[@]}; do
        echo ${ids}
      done
      echo -n "[y/n]: "
      read RESPONSE

      case $RESPONSE in
        n)
          exit 1
          ;;
        y)
          terminate_instances_with_name ${CLUSTER_NAME}
          ;;
        *)
          echo "respond with y or n."
          exit 1
      esac
    else
      echo "Invalid number of arguments"
      echo "Usage: peg down <cluster-name>"
    fi
    ;;

  ssh)
    if [[ "${nfargs}" -eq "2" ]]; then
      PARAMETER_ARR=(${parameters})
      CLUSTER_NAME=${PARAMETER_ARR[0]}
      NODE_NUM=${PARAMETER_ARR[1]}

      check_cluster_exists ${CLUSTER_NAME}
      restart_sshagent_if_needed ${CLUSTER_NAME}

      ssh -A -o "StrictHostKeyChecking no" ${REM_USER}@$(sed -n ''"$NODE_NUM"'p' ${PEG_ROOT}/tmp/${CLUSTER_NAME}/public_dns)
    else
      echo "Invalid number of arguments"
      echo "Usage: peg ssh <cluster-name> <node-number>"
    fi
    ;;

  sshcmd-node)
    if [[ "${nfargs}" -ge "3" ]]; then
      PARAMETER_ARR=(${parameters})
      CLUSTER_NAME=${PARAMETER_ARR[0]}
      NODE_NUM=${PARAMETER_ARR[1]}
      CMD=${PARAMETER_ARR[@]:2}

      check_cluster_exists ${CLUSTER_NAME}

      PUBLIC_DNS=$(sed -n ''"$NODE_NUM"'p' ${PEG_ROOT}/tmp/${CLUSTER_NAME}/public_dns)
      run_cmd_on_node ${PUBLIC_DNS} ${CMD}
    else
      echo "Invalid number of arguments"
      echo "Usage: peg sshcmd-node <cluster-name> <node-number> \"<cmd>\""
    fi
    ;;

  sshcmd-cluster)
    if [[ "${nfargs}" -ge "2" ]]; then
      PARAMETER_ARR=(${parameters})
      CLUSTER_NAME=${PARAMETER_ARR[0]}
      CMD=${PARAMETER_ARR[@]:1}

      check_cluster_exists ${CLUSTER_NAME}

      run_cmd_on_cluster ${CLUSTER_NAME} ${CMD}
    else
      echo "Invalid number of arguments"
      echo "Usage: peg sshcmd-cluster <cluster-name> \"<cmd>\""
    fi
    ;;

  scp)
    if [[ "${nfargs}" -eq "5" ]]; then
      PARAMETER_ARR=(${parameters})
      SCP_OPTION=${PARAMETER_ARR[0]}
      CLUSTER_NAME=${PARAMETER_ARR[1]}
      NODE_NUM=${PARAMETER_ARR[2]}
      LOCAL_PATH=${PARAMETER_ARR[3]}
      REM_PATH=${PARAMETER_ARR[4]}

      check_cluster_exists ${CLUSTER_NAME}

      PUBLIC_DNS=$(sed -n ''"$NODE_NUM"'p' ${PEG_ROOT}/tmp/${CLUSTER_NAME}/public_dns)
      case ${SCP_OPTION} in
        to-local|from-rem)
          scp -r -i ${PEG_ROOT}/tmp/${CLUSTER_NAME}/*.pem ${REM_USER}@${PUBLIC_DNS}:${REM_PATH} ${LOCAL_PATH}
          ;;

        to-rem|from-local)
          scp -r -i ${PEG_ROOT}/tmp/${CLUSTER_NAME}/*.pem ${LOCAL_PATH} ${REM_USER}@${PUBLIC_DNS}:${REM_PATH}
          ;;

        *)
          echo "Invalid peg scp option"
          echo "valid options:"
          echo "  to-local"
          echo "  from-local"
          echo "  to-rem"
          echo "  from-rem"
          exit 1

      esac
    else
      echo "Invalid number of arguments"
      echo "Usage: peg scp <to-local|to-rem|from-local|from-rem> <cluster-name> <node-number> <local-path> <remote-path>"
    fi
    ;;

  retag)
    if [[ "${nfargs}" -eq "2" ]]; then
      PARAMETER_ARR=(${parameters})
      CLUSTER_NAME=${PARAMETER_ARR[0]}
      NEW_CLUSTER_NAME=${PARAMETER_ARR[1]}

      retag_instance_with_name ${CLUSTER_NAME} ${NEW_CLUSTER_NAME}

      peg fetch ${NEW_CLUSTER_NAME}
    else
      echo "Invalid number of arguments"
      echo "Usage: peg retag <cluster-name> <new-cluster-name>"
    fi
    ;;

  port-forward)
    if [[ "${nfargs}" -eq "3" ]]; then
      PARAMETER_ARR=(${parameters})
      CLUSTER_NAME=${PARAMETER_ARR[0]}
      NODE_NUMBER=${PARAMETER_ARR[1]}
      PORT_CMD=${PARAMETER_ARR[2]}

      check_cluster_exists ${CLUSTER_NAME}

      port_forward ${CLUSTER_NAME} ${NODE_NUMBER} ${PORT_CMD}
    else
      echo "Invalid number of arguments"
      echo "Usage: peg port-forward <cluster-name> <node-number> <local-port>:<remote-port>"
    fi
    ;;

  start)
    if [[ "${nfargs}" -eq "1" ]]; then
      CLUSTER_NAME=${parameters}

      check_cluster_exists ${CLUSTER_NAME}

      start_instance ${CLUSTER_NAME}
    else
      echo "Invalid number of arguments"
      echo "Usage: peg start <cluster-name>"
    fi
    ;;

  stop)
    if [[ "${nfargs}" -eq "1" ]]; then
      CLUSTER_NAME=${parameters}

      check_cluster_exists ${CLUSTER_NAME}

      echo "All services should be stopped before stopping clusters."
      echo -e "Are you sure you want to stop ${color_blue}$CLUSTER_NAME${color_norm}?"
      echo -n "[y/n]: "
      read RESPONSE

      case $RESPONSE in
        n)
          exit 1
          ;;
        y)
          stop_instance ${CLUSTER_NAME}
          ;;
        *)
          echo "respond with y or n."
          exit 1
      esac
    else
      echo "Invalid number of arguments"
      echo "Usage: peg stop <cluster-name>"
    fi
    ;;

  *)
    echo "Invalid command. use -h to see available commands"
    exit 1
esac