#!/usr/libexec/platform-python
try:
    import json
except:
    import simplejson as json
import optparse
import os
import pycurl
import sys
from io import StringIO
from functools import reduce

# Exit codes (NAGIOS compliant)
EX_OK = 0
EX_WARNING = 1
EX_CRITICAL = 2
EX_UNKNOWN = 3


class FtsStalledServerProbe:
    """
    Check if the given server is stalled.
    The server is in a critical state if:
        - Server is not transferring anything
        - And the total active / number of hosts > max(number of hosts, 10)
    The server is in a warning state if:
        - It has active transfers
        - But they are less than total (active / number of hosts) * 0.8
    """

    def __init__(self, argv):
        parser = optparse.OptionParser(
            prog=os.path.basename(argv[0]),
            description=' '.join(self.__doc__.split())
        )

        parser.add_option('-H', '--host', type='str', default=None, help='FTS Host')
        parser.add_option('-m', '--monitoring', type=str, default=None,
                          help='FTS Monitoring endpoint, defaults to https://$host:8449/fts3/ftsmon/')
        parser.add_option('--cert', type=str, default=None, help='User certificate')
        parser.add_option('--key', type=str, default=None, help='User private key')

        self.options, args = parser.parse_args(argv)

        if not self.options.host:
            parser.error("-H must be specified")

        if not self.options.monitoring:
            self.options.monitoring = "https://%s:8449/fts3/ftsmon/stats/servers" % self.options.host
        else:
            self.options.monitoring += '/stats/servers'

    def __call__(self):
        buffer = StringIO()

        curl = pycurl.Curl()
        curl.setopt(pycurl.WRITEFUNCTION, buffer.write)
        curl.setopt(pycurl.FOLLOWLOCATION, True)
        curl.setopt(pycurl.URL, self.options.monitoring)
        curl.setopt(pycurl.CAPATH, '/etc/grid-security/certificates/')

        if self.options.cert:
            curl.setopt(pycurl.SSLCERT, self.options.cert)
            curl.setopt(pycurl.CAINFO, self.options.cert)

        try:
            curl.perform()
        except pycurl.error as e:
            return EX_CRITICAL, e[1]
        except Exception as e:
            return EX_CRITICAL, str(e)

        if curl.getinfo(pycurl.RESPONSE_CODE) != 200:
            return EX_CRITICAL, "Got %d" % curl.getinfo(curl.RESPONSE_CODE)

        try:
            response = json.loads(buffer.getvalue())
        except Exception as e:
            return EX_CRITICAL, 'Could not retrieve the status from the monitoring (%s)' % str(e)

        if self.options.host not in response:
            return EX_CRITICAL, 'The host is not visible in the monitoring. Maybe it is down?'

        n_hosts = len(list(response.keys()))
        total_active = reduce(lambda a, b: a + b, [h.get('active', 0) for h in list(response.values())])
        this_host_active = response[self.options.host].get('active', 0)

        msg = "Running %d out of %d" % (this_host_active, total_active)
        if this_host_active == 0 and (total_active / n_hosts) > max(n_hosts, 10):
            return EX_CRITICAL, msg
        elif this_host_active < (total_active / n_hosts) * 0.8:
            return EX_WARNING, msg

        return EX_OK, msg


if __name__ == '__main__':
    probe = FtsStalledServerProbe(sys.argv)
    (status, msg) = probe()

    if status == EX_OK:
        print(("OK -", msg))
    elif status == EX_WARNING:
        print(("WARNING -", msg))
    elif status == EX_CRITICAL:
        print(("CRITICAL -", msg))
    else:
        print(("UNKNOWN -", msg))

    sys.exit(status)
