#!/usr/bin/python
# check-es
#   Switches:
#     -x: status; status, number_of_nodes, etc., default status
#     -w (required unless -x status): warning value
#     -c (required unless -x status): critical value
#     -C (optional): comparison operator: >, >=, <, <=, ==, !=; default ==
#     -H (optional): IP/hostname; default localhost
#     -p (optional): port; default 9200
#     -h / --help (optional): help
#
#    Dependencies: ipaddress
#
#    check-es - Polls an Elasticsearch cluster node for stats
#    Copyright (C) 2017 bbb31ade7ba4103469003954b0838a96ac40b96cfc7d4c891cd34140d4849fc3
#    Copyright (C) 2018 KBase
#
#    originally from https://github.com/bbb31/check-es
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as
#    published by the Free Software Foundation, either version 3 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 Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
#

try:
    import httplib
except ImportError as e:
    import http.client as httplib

import json
import sys
import socket
import ipaddress

COMPARISON_TYPE = "str"
COMPARISON_OP   = "=="
IP_HOSTNAME     = "localhost"
IP_PORT         = 9200
STATUS_CHECK    = ""
STATUS_WARN     = ""
STATUS_CRIT     = ""
NODE_HOSTNAME   = "_local"

def check_es_help():
    print("Usage: check-es -x <status> [-w <warn> -c <crit>] [options]")
    print("Options:")
    print("\t-x: The status to check; supported values are: status, initializing_shards, number_of_data_nodes, number_of_nodes, heap_used_percent, active_shards_percent_as_number (default -x status)")
    print("\t-w: warning value (dependent upon status, required for all but -x status).")
    print("\t-c: critical value (dependent upon status, required for all but -x status).")
    print("\t-C: Comparison type: <,<=,>,>=,!=,== . Ignored if -x is status.")
    print("\t-H: Elasticsearch server to connect to. Default localhost.")
    print("\t-p: Elasticsearch server HTTP port. Default 9200.")
    print("\t-h / --help: Prints out this help document.")
    print("Version: 1.2.0")

def init_data():
    try:
      http_conn = httplib.HTTPConnection(IP_HOSTNAME, IP_PORT, timeout=5)
    # not really sure what wait for status we want here, I really just want the timeout
    # I think the status is "at least this level" so maybe red means just wait for 10s
      http_conn.request("GET", "/_cluster/health?wait_for_status=red&timeout=10s")
      health_data = json.loads(http_conn.getresponse().read())

      http_conn.request("GET", "/_nodes/%s/stats/jvm" % NODE_HOSTNAME)
      node_data = json.loads(http_conn.getresponse().read())

      http_conn.close()

      return (health_data, node_data)
    except socket.error:
      print("CRITICAL: %s is not responding." % IP_HOSTNAME)
      sys.exit(2)

def start(data):
    health_data = data[0]
    node_data = data[1]

    if STATUS_CHECK == "status":
        check_status(health_data["cluster_name"],health_data["status"])
    elif STATUS_CHECK == "heap_used_percent":
        # Verify data set just incase
        if len(node_data["nodes"]) > 0:
            node_id_key = list(node_data["nodes"].keys())[0]
            int_checker(health_data["cluster_name"],node_data["nodes"][node_id_key]["jvm"]["mem"]["heap_used_percent"])
        else:
            print("CRITICAL: Unable to retrieve nodes list from %s. Is it down?" % NODE_HOSTNAME)
            sys.exit(2)
    else:
        int_checker(health_data["cluster_name"],health_data[STATUS_CHECK])

def check_status(cluster_name, value):
    if value == STATUS_WARN:
        print("WARNING: cluster %s %s is %s." % (cluster_name, STATUS_CHECK, STATUS_WARN))
        sys.exit(1)
    elif value == STATUS_CRIT:
        print("CRITICAL: cluster %s %s is %s." % (cluster_name, STATUS_CHECK, STATUS_CRIT))
        sys.exit(2)

    print("OK: cluster %s %s is ok: %s" % (cluster_name, STATUS_CHECK,value))
    sys.exit(0)

def int_checker(cluster_name, value):

    state=3
    stateText='UNKNOWN'

    if COMPARISON_OP == "==":
        if value == int(STATUS_CRIT):
            state=2
        elif value == int(STATUS_WARN):
            state=1

    elif COMPARISON_OP == "!=":
        if value != int(STATUS_CRIT):
            state=2
        elif value != int(STATUS_WARN):
            state=1

    elif COMPARISON_OP == "<=":
        if value <= int(STATUS_CRIT):
            state=2
        elif value <= int(STATUS_WARN):
            state=1

    elif COMPARISON_OP == "<":
        if value < int(STATUS_CRIT):
            state=2
        elif value < int(STATUS_WARN):
            state=1

    elif COMPARISON_OP == ">":
        if value > int(STATUS_CRIT):
            state=2
        elif value > int(STATUS_WARN):
            state=1

    elif COMPARISON_OP == ">=":
        if value >= int(STATUS_CRIT):
            state=2
        elif value >= int(STATUS_WARN):
            state=1

    if state == 3:
        state=0
        stateText=("OK: cluster %s %s is ok: %s" % (cluster_name, STATUS_CHECK,value))
        if STATUS_CHECK == "heap_used_percent":
            stateText=("OK: node %s is ok: %s" % (STATUS_CHECK,value))
    if state == 2:
        stateText=("CRITICAL: cluster %s %s is %d which is %s %s" % (cluster_name, STATUS_CHECK, value, COMPARISON_OP, STATUS_CRIT))
        if STATUS_CHECK == "heap_used_percent":
            stateText=("CRITICAL: node %s is %d which is %s %s" % (STATUS_CHECK, value, COMPARISON_OP, STATUS_CRIT))
    if state == 1:
        stateText=("WARNING: cluster %s %s is %d which is %s %s" % (cluster_name, STATUS_CHECK, value, COMPARISON_OP, STATUS_WARN))
        if STATUS_CHECK == "heap_used_percent":
            stateText=("WARNING: node %s is %d which is %s %s" % (STATUS_CHECK, value, COMPARISON_OP, STATUS_WARN))

    print (stateText)
    sys.exit(state)

# help
if "--help" in sys.argv or "-h" in sys.argv:
    check_es_help()
    sys.exit(0)

tmp_args = sys.argv[1:]
tmp_args = dict([(tmp_args[i],tmp_args[i+1]) for i,b in enumerate(tmp_args[1:]) if i % 2 == 0])

# setup check
if "-x" not in sys.argv:
    tmp_args["-x"] = "status"

# setup -x
if tmp_args["-x"] == "status" or tmp_args["-x"] == "initializing_shards" or tmp_args["-x"] == "number_of_data_nodes" or tmp_args["-x"] == "number_of_nodes" or tmp_args["-x"] == "heap_used_percent" or tmp_args["-x"] == "active_shards_percent_as_number":
    STATUS_CHECK = tmp_args["-x"]

    if tmp_args["-x"] == "status":
        COMPARISON_TYPE = "str"
    elif tmp_args["-x"] == "initializing_shards":
        COMPARISON_TYPE = "int"
    elif tmp_args["-x"] == "number_of_data_nodes":
        COMPARISON_TYPE = "int"
    elif tmp_args["-x"] == "number_of_nodes":
        COMPARISON_TYPE = "int"
    elif tmp_args["-x"] == "heap_used_percent":
        NODE_MODE = True
        COMPARISON_TYPE = "int"
    elif tmp_args["-x"] == "active_shards_percent_as_number":
        COMPARISON_TYPE = "int"

else:
    print("%s is not supported; see help." % (tmp_args["-x"]) )
    sys.exit(3)

# setup -w
try:
    if STATUS_CHECK == "status":
        STATUS_WARN = 'yellow'
    elif COMPARISON_TYPE == "int":
        if "-w" not in sys.argv:
            print("-w required")
            sys.exit(3)
        try:
            x_val = int(tmp_args["-w"])

            if x_val < 0:
                print("Negative values? Really?")
                sys.exit(3)
            else:
                STATUS_WARN = x_val
        except ValueError:
                print("%s is not valid." % tmp_args["-w"])
                sys.exit(3)

except Exception as e:
    print(e)

# setup -c
try:
    if STATUS_CHECK == "status":
        STATUS_CRIT = 'red'
    elif COMPARISON_TYPE == "int":
        if "-c" not in sys.argv:
            print("-c required")
            sys.exit(3)
        try:
            x_val = int(tmp_args["-c"])

            if x_val < 0:
                print("Negative values? Really?")
                sys.exit(3)
            else:
                STATUS_CRIT = x_val
        except ValueError:
            print("%s is not valid." % tmp_args["-c"])
            sys.exit(3)
except Exception as e:
    print(e)


# setup -C
if "-C" in tmp_args:
    if tmp_args["-C"] == "<=" or tmp_args["-C"] == "<" or tmp_args["-C"] == ">" or tmp_args["-C"] == ">=" or tmp_args["-C"] == "==" or tmp_args["-C"] == "!=":
        COMPARISON_OP = tmp_args["-C"]
    else:
        print("%s is an invalid comparison operator; only >,>=,<,<=,==,!= are supported")
        sys.exit(3)

# setup -H
if "-H" in tmp_args:
    bad_host_count = 0

    # check if IP
    try:
        ipaddress.ip_address(tmp_args["-H"].decode("utf-8"))
    except ValueError:
        bad_host_count += 1

    # check if host
    try:
        socket.gethostbyname(tmp_args["-H"])
    except socket.gaierror:
        bad_host_count += 1

    if bad_host_count > 1:
        print("Invalid IP/hostname %s" % tmp_args["-H"])
        sys.exit(3)

    IP_HOSTNAME = tmp_args["-H"]
    NODE_HOSTNAME = IP_HOSTNAME

# setup -p
if "-p" in tmp_args:
    try:
        int(tmp_args["-p"])

        if int(tmp_args["-p"]) < 1 or int(tmp_args["-p"]) > 65535:
            print("%s is not a valid port." % tmp_args["-p"])
            sys.exit(3)

        IP_PORT = int(tmp_args["-p"])
    except ValueError:
        print("%s is not a valid port." % tmp_args["-p"])
        sys.exit(3)

data = init_data()
start(data)
