check_doh

#!/usr/bin/env python3
#-------------------------------------------------------------------------------

import sys
import os
import argparse
import socket
import httpx
import dns.message
import dns.query
import dns.rdatatype
import time

#-------------------------------------------------------------------------------

_candebug = False
def candebug():
    global _candebug
    return _candebug
def setcandebug(value):
    global _candebug
    _candebug = value

def infomsg(msg):
    if candebug() == True:
        print(msg, flush=True)

def exitnagios(status,message):
    if status=="OK":
        exitcode = 0
    elif status=="WARNING":
        exitcode = 1
    elif status=="CRITICAL":
        exitcode = 2
    elif status=="UNKNOWN":
        exitcode = 3
    else:
        exitcode = 4
    print(status+": "+message, flush=True)
    sys.exit(exitcode)

#-------------------------------------------------------------------------------

original_getaddrinfo = socket.getaddrinfo

def forced_ipv6_gai_family(*args, **kwargs):
    global original_getaddrinfo
    responses = original_getaddrinfo(*args, **kwargs)
    return [response
            for response in responses
            if response[0] == socket.AF_INET6]

def forced_ipv4_gai_family(*args, **kwargs):
    global original_getaddrinfo
    responses = original_getaddrinfo(*args, **kwargs)
    return [response
            for response in responses
            if response[0] == socket.AF_INET]

#-------------------------------------------------------------------------------

def makedohtest(retries, targetdohurl, testquery):
    result = False
    infomsg("preparing client")
    timeout = httpx.Timeout(5.0)

    counter = 0
    while True:
        client = httpx.Client(timeout=timeout)
        try:
            q = dns.message.make_query(testquery, dns.rdatatype.ANY)
            infomsg("query ready")
            r = dns.query.https(q, targetdohurl, session=client, timeout=5.0)
            infomsg("got answer!")
            if len(r.answer)>0:
                result = True
        except Exception as e:
            exitnagios("CRITICAL","url "+targetdohurl+" caused an exception "+str(e))
        finally:
            infomsg("ended")
            if client != None:
                client.close()

        counter = counter+1
        if result==True:
            break
        if counter>retries:
            infomsg("aborted!")
            break
        else:
            infomsg("sleeping before retry!")
            time.sleep(5)

    if result == True:
        exitnagios("OK","DoH for "+targetdohurl+" is working!")
    else:
        exitnagios("CRITICAL","DoH for "+targetdohurl+" is broken")

#-------------------------------------------------------------------------------

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("-4", "--ipv4",              action="store_true",  dest="ipv4",             default=False,                                help="use ipv4 (exclusive with ipv6)")
    parser.add_argument("-6", "--ipv6",              action="store_true",  dest="ipv6",             default=False,                                help="use ipv6 (exclusive with ipv4)")
    parser.add_argument("-0", "--ipv0",              action="store_true",  dest="ipv0",             default=False,                                help="do not specify ip protocol version")
    parser.add_argument("-t", "--targetdohurl",      action="store",                                                                              help="target url for the DoH server")
    parser.add_argument("-q", "--testquery",         action="store",                                default="www.google.com.",                    help="query to resolve")
    parser.add_argument("-®", "--debug",             action="store_true",  dest="debug",            default=False,                                help="be more verbose")
    args = parser.parse_args()
    return args

def main():
    args = parse_args()

    setcandebug(args.debug)

    protocol = None
    exclist = 0
    if args.ipv0 == True :
        protocol = ""
        exclist = exclist+1
    if args.ipv4 == True :
        protocol = "-4"
        exclist = exclist+1
        socket.getaddrinfo = forced_ipv4_gai_family
    if args.ipv6 == True :
        protocol = "-6"
        exclist = exclist+1
        socket.getaddrinfo = forced_ipv6_gai_family
    if (exclist>1) :
        exitnagios("CRITICAL","select ip protocol version 4 or 6, or 0 to relay on OS handling")
    elif (exclist<1) :
        exitnagios("CRITICAL","no protocol selected")

    if not args.targetdohurl:
        exitnagios("CRITICAL","Must give a targetdohurl.")

    retries=3
    makedohtest(retries,args.targetdohurl,args.testquery)

if __name__ == "__main__":
    main()

#-------------------------------------------------------------------------------