#!/usr/bin/env python3

"""
DNS Health Check

Description:
    - Performs a DNS resolution test for a random or specified hostname.
    - Uses local resolvers from /etc/resolv.conf, with support for optional additional resolvers via --resolver.
    - Expects at least 2 local resolvers to be configured, or returns CRITICAL.
    - Returns:
        0 = OK (all resolvers succeeded)
        2 = CRITICAL (any resolver failed or too few local resolvers)
        3 = UNKNOWN (no resolvers found)

Usage examples:
    check_resolver.py
    check_resolver.py --hostname www.example.com
    check_resolver.py --resolver 1.1.1.1 --resolver 8.8.8.8
"""

import subprocess
import ipaddress
import random
import argparse
import sys

HOSTNAME = "wwwmi3-ss1.a2hosting.com"

def load_local_resolvers():
    resolvers = []
    try:
        with open("/etc/resolv.conf", "r") as f:
            for line in f:
                if line.startswith("nameserver"):
                    parts = line.strip().split()
                    if len(parts) == 2:
                        resolvers.append(parts[1])
    except Exception:
        pass
    return resolvers

def is_valid_ip(ip):
    try:
        ipaddress.ip_address(ip)
        return True
    except ValueError:
        return False

def run_dig(server, hostname):
    try:
        proc = subprocess.Popen(
            ["dig", f"@{server}", hostname, "+short"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        stdout, stderr = proc.communicate(timeout=5)
        output = stdout.decode().strip().splitlines()

        for line in output:
            if is_valid_ip(line):
                return True

        return False

    except subprocess.TimeoutExpired:
        return False
    except Exception as e:
        return False

def main():
    parser = argparse.ArgumentParser(description="DNS Health Check")
    parser.add_argument("--hostname", help="Hostname to resolve (random if not set)")
    parser.add_argument("--resolver", action="append", help="Additional resolver(s) to use", default=[])
    args = parser.parse_args()

    # Load and validate local resolvers
    local_resolvers = load_local_resolvers()
    if len(local_resolvers) < 2:
        print(f"CRITICAL - Only {len(local_resolvers)} local resolver{'s' if len(local_resolvers) != 1 else ''} found in /etc/resolv.conf")
        sys.exit(2)

    # Generate random hostname
    if args.hostname:
        hostname = args.hostname
    else:
        rand = random.randint(1, 100000)
        hostname = f"{rand}.{HOSTNAME}"

    # Combine resolvers
    resolvers = list(dict.fromkeys(local_resolvers + args.resolver))

    if not resolvers:
        print("UNKNOWN - No resolvers found or specified.")
        sys.exit(3)

    # Perform DNS checks
    failed = []
    for server in resolvers:
        if not run_dig(server, hostname):
            failed.append(server)

    total = len(resolvers)
    failed_count = len(failed)

    if failed_count == 0:
        print(f"OK - All {total} resolvers successfully resolved {hostname}")
        sys.exit(0)
    else:
        plural = "is" if failed_count == 1 else "are"
        print(f"CRITICAL - {failed_count} of {total} resolvers ({', '.join(failed)}) {plural} failing to resolve: {hostname}")
        sys.exit(2)

if __name__ == "__main__":
    main()