#!/bin/sh

#   This program is free software; you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation; either version 2 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

PROGNAME=`basename $0`
VERSION="Version 0.3.0,"
AUTHOR="Andrew Lyon, Based on Mike Adolphs (http://www.matejunkie.com/) check_nginx.sh code. Authentication support added by Ryan Gallant"

ST_OK=0
ST_WR=1
ST_CR=2
ST_UK=3
epoch=`date +%s`
hostname="localhost"
user=elastic
pass=changeme
authentication="False"
port=9200
status_page="_cluster/health"
output_dir=/tmp
scheme=http
carbon_key="system.:::name:::.cluster.app.elasticsearch.cluster"



print_version() {
    echo "$VERSION $AUTHOR"
}

print_help() {
    print_version $PROGNAME $VERSION
    echo ""
    echo "$PROGNAME is a Nagios plugin to check the cluster status of elasticsearch."
    echo "It also parses the status page to get a few useful variables out, and return"
    echo "them in the output."
    echo ""
    echo "$PROGNAME -H localhost -P 9200 -o /tmp"
    echo ""
    echo "Options:"
    echo "  -H/--hostname)"
    echo "     Defines the hostname. Default is: localhost"
    echo "  -P/--port)"
    echo "     Defines the port. Default is: 9200"
    echo "  -s/--secure)"
    echo "     Use TLS. Defaults to false."
    echo "  -o/--output-directory)"
    echo "     Specifies where to write the tmp-file that the check creates."
    echo "     Default is: /tmp"
    echo "  -u/--username)"
    echo "    Username for elasticsearch. Turns on authentication mode when set"
    echo "  -p/--password)"
    echo "    Password for elasticsearch. Turns on authentication mode when set"
    echo "  -a/--auth)"
    echo "    Turns on authentication mode with default credentials." 
    exit $ST_UK
}

while test -n "$1"; do
    case "$1" in
        -help|-h)
            print_help
            exit $ST_UK
            ;;
        --version|-v)
            print_version $PROGNAME $VERSION
            exit $ST_UK
            ;;
        --hostname|-H)
            hostname=$2
            shift
            ;;
        --secure|-s)
            scheme=https
            ;;
        --port|-P)
            port=$2
            shift
            ;;
        --password|-p)
            pass=$2
            authentication="True"
            shift
            ;;
        --username|-u)
            user=$2
            authentication="True"
            shift
            ;;
        --auth|-a)
            authentication="True"
            ;;
        --output-directory|-o)
            output_dir=$2
            shift
            ;;
        --carbon-server|-c)
            carbon_server=$2
            shift
            ;;
        --carbon-port|-C)
            carbon_port=$2
            shift
            ;;
        --carbon-key|-k)
            carbon_key=$2
            shift
            ;;
        *)
        echo "Unknown argument: $1"
        print_help
        exit $ST_UK
        ;;
    esac
    shift
done

if [ "$authentication" = "True" ]; then 
    pass="--password=${pass}"
    user="--user=${user}"
fi 

get_status() {
    filename=${PROGNAME}-${hostname}-${status_page}.1
    filename=`echo $filename | tr -d '\/'`
    filename=${output_dir}/${filename}
    wget -q -t 3 -T 3 ${user} ${pass} $scheme://${hostname}:${port}/${status_page}?pretty=true -O ${filename}
}

get_vals() {
    name=$(grep '\"cluster_name\" :' ${filename} | awk -F '"' '{print $4}')
    status=$(grep '\"status\" :' ${filename} | awk -F '"' '{print $4}')
    timed_out=`grep \"timed_out\" ${filename} | awk '{print $3}' | sed 's|[\r",]||g'`
    number_nodes=`grep \"number_of_nodes\" ${filename} | awk '{print $3}' | sed 's|[\r",]||g'`
    number_data_nodes=`grep \"number_of_data_nodes\" ${filename} | awk '{print $3}' | sed 's|[\r",]||g'`
    active_primary_shards=`grep \"active_primary_shards\" ${filename} | awk '{print $3}' | sed 's|[\r",]||g'`
    active_shards=`grep \"active_shards\" ${filename} | awk '{print $3}' | sed 's|[\r",]||g'`
    relocating_shards=`grep \"relocating_shards\" ${filename} | awk '{print $3}' | sed 's|[\r",]||g'`
    initializing_shards=`grep \"initializing_shards\" ${filename} | awk '{print $3}' | sed 's|[\r",]||g'`
    delayed_unassigned_shards=`grep \"delayed_unassigned_shards\" ${filename} | awk '{print $3}' | sed 's|[\r",]||g'`
    unassigned_shards=`grep -w '\"unassigned_shards\"' ${filename} | awk '{print $3}' | sed 's|[\r",]||g'`
    rm -f ${filename}

    # Determine the Nagios Status and Exit Code
    if [ "$status" = "red" ]; then
        NAGSTATUS="CRITICAL"
        EXST=$ST_CR
    elif [ "$status" = "yellow" ]; then
        NAGSTATUS="WARNING"
        EXST=$ST_WR
    elif [ "$status" = "green" ]; then
        NAGSTATUS="OK"
        EXST=$ST_OK
    else
        NAGSTATUS="UNKNOWN"
        EXST=$ST_UK
    fi
}

do_output() {
    output="elasticsearch ($name) is running. \
status: $status; \
timed_out: $timed_out; \
number_of_nodes: $number_nodes; \
number_of_data_nodes: $number_data_nodes; \
active_primary_shards: $active_primary_shards; \
active_shards: $active_shards; \
relocating_shards: $relocating_shards; \
initializing_shards: $initializing_shards; \
delayed_unassigned_shards: $delayed_unassigned_shards; \
unassigned_shards: $unassigned_shards "
}

do_perfdata() {
    #perfdata="'idle'=$iproc 'active'=$aproc 'total'=$tproc"
    perfdata="'active_primary'=$active_primary_shards 'active'=$active_shards 'relocating'=$relocating_shards 'init'=$initializing_shards 'delay_unass'=$delayed_unassigned_shards 'unass'=$unassigned_shards"
}

do_graphite() {
    if [ "$carbon_server" != "" -a "$carbon_port" != "" ]; then
        key=$(echo $carbon_key | sed "s/:::name:::/$name/")
        echo "$key.cluster.status                $EXST                  $epoch" | nc -w 2 $carbon_server $carbon_port
        echo "$key.cluster.nodes.ttl             $number_nodes          $epoch" | nc -w 2 $carbon_server $carbon_port
        echo "$key.cluster.nodes.data            $number_data_nodes     $epoch" | nc -w 2 $carbon_server $carbon_port
        echo "$key.cluster.shards.active         $active_shards         $epoch" | nc -w 2 $carbon_server $carbon_port
        echo "$key.cluster.shards.active_primary $active_primary_shards $epoch" | nc -w 2 $carbon_server $carbon_port
        echo "$key.cluster.shards.initializing   $initializing_shards   $epoch" | nc -w 2 $carbon_server $carbon_port
        echo "$key.cluster.shards.relocating     $relocating_shards     $epoch" | nc -w 2 $carbon_server $carbon_port
        echo "$key.cluster.shards.delayed_unassigned     $delayed_unassigned_shards     $epoch" | nc -w 2 $carbon_server $carbon_port
        echo "$key.cluster.shards.unassigned     $unassigned_shards     $epoch" | nc -w 2 $carbon_server $carbon_port
        unset key
    fi
}

# Here we go!
which wget >/dev/null 2>&1
if [ "$?" != "0" ]; then
    echo "CRITICAL - wget is not installed"
    exit $ST_CR
fi
get_status
if [ ! -s "$filename" ]; then
    echo "CRITICAL - Could not connect to server $hostname"
    exit $ST_CR
else
    get_vals
    if [ -z "$name" ]; then
        echo "CRITICAL - Error parsing server output"
        exit $ST_CR
    else
        do_output
        do_perfdata
    do_graphite
    fi
fi

COMPARE=$listql

echo "${NAGSTATUS} - ${output} | ${perfdata}"
exit $EXST
