#!/usr/bin/env python3
import argparse
import socket
import sys
import time

import dns.message
import dns.query


def resolve_dns_server_ip(server):
    try:
        ip_address = socket.getaddrinfo(server, None, family=socket.AF_INET)[0][4][0]
    except socket.gaierror:
        raise ValueError("Failed to resolve DNS server domain")
    return ip_address


def main():
    parser = argparse.ArgumentParser(description='Check DNS-over-TLS connectivity and response time')
    parser.add_argument('-c', '--critical', type=float, default=500, help='Critical threshold for response time in milliseconds (default: 500)')
    parser.add_argument('-w', '--warning', type=float, default=300, help='Warning threshold for response time in milliseconds (default: 300)')
    parser.add_argument('-d', '--domain', type=str, default='google.com', help='Domain to resolve (default: google.com)')
    parser.add_argument('-s', '--server', type=str, required=True, help='DNS server address')
    parser.add_argument('-p', '--port', type=int, default=853, help='Port to use for DNS-over-TLS (default: 853)')
    args = parser.parse_args()

    try:
        server_ip = resolve_dns_server_ip(args.server)
    except ValueError as e:
        print(f'CRITICAL - {str(e)}')
        sys.exit(2)

    error_message = None
    query = dns.message.make_query(args.domain, dns.rdatatype.A)
    start_time = time.time()
    try:
        response = dns.query.tls(query, server_ip, port=args.port, timeout=args.critical)
    except dns.exception.SyntaxError:
        error_message = "invalid IP address or domain name"
    except dns.exception.Timeout:
        error_message = "the DNS operation timed out"
    except Exception as e:
        error_message = e

    if error_message:
        print(f'CRITICAL - DNS-over-TLS query failed: {error_message}')
        sys.exit(2)

    response_time = time.time() - start_time
    response_time_ms = int(response_time * 1000)

    if response.rcode() != dns.rcode.NOERROR:
        print(f'CRITICAL - DNS query returned error: {dns.rcode.to_text(response.rcode())}')
        sys.exit(2)

    resolved_ips = [str(rr.address) for rr in response.answer[0] if rr.rdtype == dns.rdatatype.A]
    perfdata = f'Resolved IPs: {", ".join(resolved_ips)} | response_time={response_time_ms}ms;{args.warning};{args.critical};0;'

    if response_time_ms >= args.critical:
        print(f'CRITICAL - DNS-over-TLS response time: {response_time:.2f} sec.')
        sys.exit(2)
    elif response_time_ms >= args.warning:
        print(f'WARNING - DNS-over-TLS response time: {response_time:.2f} sec.')
        sys.exit(1)
    else:
        print(f'OK - DNS-over-TLS response time: {response_time:.2f} sec. {perfdata}')
        sys.exit(0)


if __name__ == '__main__':
    main()