#!/usr/bin/env python3
import argparse, sys, base64, time, warnings
import dns.message, dns.rdatatype, dns.rcode
import httpx
warnings.filterwarnings("ignore")

def perform_doh_query(server_url, name, qtype, insecure, method, timeout):
    try:
        q = dns.message.make_query(name, getattr(dns.rdatatype, qtype))
        wire = q.to_wire()
        headers = {'accept': 'application/dns-message'}
        verify = not insecure

        t0 = time.time()
        with httpx.Client(http2=True, verify=verify, timeout=timeout) as c:
            if method == 'GET':
                dns_param = base64.urlsafe_b64encode(wire).decode().rstrip('=')
                r = c.get(server_url, params={'dns': dns_param}, headers=headers)
            else:
                headers['content-type'] = 'application/dns-message'
                r = c.post(server_url, content=wire, headers=headers)
        r.raise_for_status()
        rt_ms = int((time.time() - t0) * 1000)
        return dns.message.from_wire(r.content), rt_ms
    except Exception as e:
        print(f"CRITICAL: Request failed - {e}")
        sys.exit(2)

def check_nagios_status(resp, name, qtype, rt_ms):
    if resp.rcode() != dns.rcode.NOERROR:
        print(f"CRITICAL: DNS error - {dns.rcode.to_text(resp.rcode())} | rtt_ms={rt_ms}")
        sys.exit(2)
    if not resp.answer or len(resp.answer[0]) == 0:
        print(f"WARNING: No {qtype} records found for {name} | rtt_ms={rt_ms} records=0")
        sys.exit(1)
    cnt = len(resp.answer[0])
    print(f"OK: Found {cnt} {qtype} records for {name} | rtt_ms={rt_ms} records={cnt}")
    sys.exit(0)

if __name__ == '__main__':
    p = argparse.ArgumentParser(description='Nagios check for DNS-over-HTTPS (RFC 8484, HTTP/2)')
    p.add_argument('--server', required=True, help='DoH URL, np. https://host/dns-query')
    p.add_argument('--name', required=True, help='FQDN do sprawdzenia')
    p.add_argument('--type', default='A', help='Typ rekordu (A, AAAA, TXT, itp.)')
    p.add_argument('--insecure', action='store_true', help='Pomiń weryfikację TLS')
    p.add_argument('--get', action='store_true', help='Użyj GET zamiast POST')
    p.add_argument('--timeout', type=float, default=5.0, help='Timeout w sekundach (domyślnie 5)')
    args = p.parse_args()

    method = 'GET' if args.get else 'POST'
    resp, rt = perform_doh_query(args.server, args.name, args.type, args.insecure, method, args.timeout)
    check_nagios_status(resp, args.name, args.type, rt)