#!/usr/bin/env python
import getopt
import os
import pycurl
import sys

from StringIO import StringIO
from xml.dom import minidom

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

# Message
# Hardcoded to avoid dependencies
SOAP_VERSION_REQUEST = """
<?xml version="1.0" encoding="UTF-8"?>
<SOAP-ENV:Envelope
    xmlns:SOAP-ENC="http://schemas.xmlsoap.org/soap/encoding/"
    xmlns:ns0="http://schemas.xmlsoap.org/soap/encoding/"
    xmlns:ns1="http://schemas.xmlsoap.org/soap/envelope/"
    xmlns:ns2="http://glite.org/wsdl/services/org.glite.data.transfer.fts"
    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
    xmlns:SOAP-ENV="http://schemas.xmlsoap.org/soap/envelope/"
    SOAP-ENV:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/">
   <SOAP-ENV:Header/>
   <ns1:Body>
      <ns2:getInterfaceVersion/>
   </ns1:Body>
</SOAP-ENV:Envelope>
"""

class FtsProbe:
    """
    Checks if the FTS server is up and running using the
    getVersion call.
    """

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

        # Defaults
        self.abortMsg = None
        self.help     = False
        self.host     = None
        self.port     = 8443
        self.cert     = None
        self.key      = None

        # Process args
        try:
            opts, args = getopt.getopt(argv[1:], 'hH:p:', ['help', 'host=', 'port=', 'cert=', 'key='])
            for opt, arg in opts:
                if opt in ('-h', '--help'):
                    self.help = True
                elif opt in ('-H', '--host'):
                    self.host = arg
                elif opt in ('-p', '--port'):
                    self.port = int(arg)
                elif opt == '--cert':
                    self.cert = arg
        except Exception, e:
            self.abortMsg = str(e)

        if not self.abortMsg and self.host is None:
            self.abortMsg = "-H must be specified"

    def usage(self):
        return \
"""%s

Usage:
    %s -H <host> [-p <port>] [--cert <user-cert>] [--key <user-key>]

Options:
    -H, --host Specifies the host to check
    -p, --port Specifies the FTS port. Default: 8443
        --cert X509 certificate to use
""" % (self.description, self.scriptName)


    def __call__(self):
        if self.help:
            return (EX_OK, self.usage())
        elif self.abortMsg:
            return (EX_CRITICAL, self.abortMsg)

        curl = pycurl.Curl()
        curl.setopt(pycurl.URL, "https://%s:%s" % (self.host, self.port))
        curl.setopt(pycurl.CAPATH, '/etc/grid-security/certificates')
        if self.cert:
            curl.setopt(pycurl.CAINFO, self.cert)
            curl.setopt(pycurl.SSLCERT, self.cert)

        curl.setopt(pycurl.POST, True)
        curl.setopt(pycurl.HTTPHEADER, ['Content-Type: text/xml'])
        curl.setopt(pycurl.POSTFIELDS, SOAP_VERSION_REQUEST)

        buffer = StringIO()
        curl.setopt(pycurl.WRITEFUNCTION, buffer.write)

        try:
            curl.perform()
        except Exception, e:
            return (EX_CRITICAL, str(e))
       
        if curl.getinfo(pycurl.RESPONSE_CODE) != 200:
            return (EX_CRITICAL, "Got %d" % curl.getinfo(pycurl.RESPONSE_CODE))
 
        try:
            doc = minidom.parseString(buffer.getvalue())
            response = doc.getElementsByTagName('getInterfaceVersionReturn')[0]
            version = None
            for child in response.childNodes:
                if child.nodeType == child.TEXT_NODE:
                    version = child.data
            if version is None:
                raise Exception('Could not find the version')
        except Exception, e:
            return (EX_CRITICAL, "Response is malformed (%s)" % str(e))

        return (EX_OK, "FTS %s" % version)


if __name__ == '__main__':
    probe = FtsProbe(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)

