#!/usr/bin/env python3

"""Monitoring plugin (Nagios-compatible) for watching a DoH resolver.

The monitoring plugin API is documented at
<https://www.monitoring-plugins.org/doc/guidelines.html>.

"""

# Do not touch
# https://www.monitoring-plugins.org/doc/guidelines.html#AEN78
STATE_OK = 0
STATE_WARNING = 1
STATE_CRITICAL = 2
STATE_UNKNOWN = 3
STATE_DEPENDENT = 4

# http://pycurl.io/docs/latest
import pycurl

# http://www.dnspython.org/
import dns.message

import io
import sys
import base64
import urllib.parse
import socket
import re
import getopt

host = None
vhostname = None
path = None
lookup = None
ltype = 'AAAA'
post = False
insecure = False
head = False
# TODO add an option: a string which is expected in the DNS response

try:
    optlist, args = getopt.getopt (sys.argv[1:], "H:n:p:V:t:Pih")
    for option, value in optlist:
        if option == "-H":
            host = value
        elif option == "-V":
            vhostname = value
        elif option == "-n":
            lookup = value
        elif option == "-t":
            ltype = value
        elif option == "-p":
            path = value
        elif option == "-P":
            post = True
        elif option == "-i":
            insecure = True
        elif option == "-h":
            head = True
        else:
            # Should never occur, it is trapped by getopt
            print("Unknown option %s" % option)
            sys.exit(STATE_UNKNOWN)
except getopt.error as reason:
    print("Option parsing problem %s" % reason)
    sys.exit(STATE_UNKNOWN)
if len(args) > 0:
    print("Too many arguments (\"%s\")" % args)
    sys.exit(STATE_UNKNOWN)
if host is None or lookup is None:
    print("Host and name to lookup are necessary")
    sys.exit(STATE_UNKNOWN)
if post and head:
    print("POST or HEAD but not both")
    sys.exit(STATE_UNKNOWN)
    
if vhostname is not None:
    url = "https://%s/" % vhostname # host is ignored in that case, which is a bit strange
else:
    url = "https://%s/" % host
if path is not None:
    url += path
                   
try:
    buffer = io.BytesIO()
    c = pycurl.Curl()
    message = dns.message.make_query(lookup, dns.rdatatype.from_text(ltype))
    message.id = 0 # DoH requests that
    if head:
        c.setopt(pycurl.NOBODY, True)
    if post:
        c.setopt(c.URL, url)
        data = message.to_wire()
        c.setopt(pycurl.POST, True)
        c.setopt(pycurl.POSTFIELDS, data)
    else:
        dns_req = base64.urlsafe_b64encode(message.to_wire()).decode('UTF8').rstrip('=')
        c.setopt(c.URL, url + ("?dns=%s" % dns_req))
    c.setopt(pycurl.HTTPHEADER, ["Content-type: application/dns-message"])
    c.setopt(c.WRITEDATA, buffer)
    if insecure:
        c.setopt(pycurl.SSL_VERIFYPEER, False)   
        c.setopt(pycurl.SSL_VERIFYHOST, False)
    # Does not work if pycurl was not compiled with nghttp2 (recent Debian
    # packages are OK) https://github.com/pycurl/pycurl/issues/477
    c.setopt(pycurl.HTTP_VERSION, pycurl.CURL_HTTP_VERSION_2)
    try:
        c.perform()
    except Exception as e:
        print("%s ERROR - %s" % (url, "Cannot connect: \"%s\"" % e))
        sys.exit(STATE_CRITICAL)
    rcode = c.getinfo(pycurl.RESPONSE_CODE)
    c.close()
    if rcode == 200:
        if not head:
            body = buffer.getvalue()
            try:
                response = dns.message.from_wire(body) # May be we should test the DNS response code as well?
            except dns.name.BadLabelType as e:
                print("%s ERROR - %s" % (url, "Not a DNS reply, is it a DoH server? \"%s\"" % e))
                sys.exit(STATE_CRITICAL)
            print("%s OK - %s" % (url, "No error for %s/%s, %i bytes received" % (lookup, ltype, sys.getsizeof(body))))
        else:
            print("%s OK - %s" % (url, "No error"))
        sys.exit(STATE_OK)
    else:
        body =  buffer.getvalue()
        if len(body) == 0:
            body = b"[No details]"
        print("%s HTTP error - %i: %s" % (url, rcode, body.decode()))
        sys.exit(STATE_CRITICAL)
except Exception as e:
    print("%s UNKNOWN - %s" % (url, "Unknown internal error: \"%s\"" % e))
    sys.exit(STATE_UNKNOWN)