#!/usr/bin/env python3
"""Read-only, fail-closed public-DNS preflight for the approved SRS DNS profile.

Checks the authoritative nameservers, bypassing local hosts/systemd-resolved.
This deliberately validates only the approved single-host DNS profile:
exactly one MX pointing to this host and one SPF record `v=spf1 mx -all`.
It is NOT a general-purpose SPF evaluator.
"""
import argparse
import ipaddress
import re
import shutil
import subprocess
import sys

# A public resolver is used only to discover the delegated authoritative NS and
# their numeric IPs. The actual TXT/MX/A/AAAA values are read from each NS.
BOOTSTRAP_RESOLVERS = ('1.1.1.1', '9.9.9.9')
DNS_TIMEOUT = 5


def dig(server, name, kind, authoritative=False):
    command = ['dig', '@' + server, '+time=2', '+tries=1', '+noall', '+comments', '+answer']
    if authoritative:
        command.append('+norecurse')
    command.extend([name, kind])
    try:
        proc = subprocess.run(command, capture_output=True, text=True, timeout=DNS_TIMEOUT, check=False)
    except (OSError, subprocess.TimeoutExpired) as exc:
        raise ValueError(f'DNS request failed ({server}, {name}, {kind}): {exc}') from exc
    if proc.returncode != 0:
        raise ValueError(f'DNS request failed ({server}, {name}, {kind}): {proc.stderr.strip() or proc.returncode}')
    header = re.search(r'HEADER.*?status:\s*(\w+)', proc.stdout)
    if not header or header.group(1) != 'NOERROR':
        raise ValueError(f'DNS returned {header.group(1) if header else "no status"}: {server} {name} {kind}')
    flags = re.search(r'^;; flags:\s*(.*?);\s*QUERY:', proc.stdout, re.M)
    if authoritative and (not flags or 'aa' not in flags.group(1).split()):
        raise ValueError(f'Authoritative response missing AA flag: {server} {name} {kind}')
    rows = []
    for line in proc.stdout.splitlines():
        parts = line.split(None, 4)
        if len(parts) != 5 or parts[2].upper() != 'IN':
            continue
        owner, _ttl, _cls, rr_type, data = parts
        if rr_type.upper() == kind.upper() and owner.rstrip('.').lower() == name.rstrip('.').lower():
            rows.append(data.strip())
    return rows


def public_lookup(name, kind):
    failures = []
    for server in BOOTSTRAP_RESOLVERS:
        try:
            return dig(server, name, kind)
        except ValueError as exc:
            failures.append(str(exc))
    raise ValueError('Public NS bootstrap failed: ' + ' | '.join(failures))


def discover_nameservers(domain):
    labels = domain.rstrip('.').split('.')
    # Pick closest delegated zone. A zone without its own delegation has no NS.
    for i in range(len(labels) - 1):
        zone = '.'.join(labels[i:])
        ns = public_lookup(zone, 'NS')
        if ns:
            names = sorted({n.rstrip('.').lower() for n in ns})
            return zone, names
    raise ValueError(f'Cannot discover authoritative nameservers for {domain}')


def addr_of_ns(ns):
    for kind in ('A', 'AAAA'):
        for value in public_lookup(ns, kind):
            try:
                return str(ipaddress.ip_address(value))
            except ValueError:
                continue
    raise ValueError(f'Cannot resolve nameserver {ns} to a public IP address')


def parse_txt(value):
    # DNS may split a single TXT resource record into multiple quoted strings.
    parts = re.findall(r'"((?:\\.|[^"\\])*)"', value)
    if not parts:
        raise ValueError('Unexpected TXT presentation format: ' + value)
    return ''.join(parts).replace('\\"', '"')


def check_server(ns_name, ns_ip, domain, hostname, bind4, bind6):
    errors = []
    def query(name, kind):
        try:
            return dig(ns_ip, name, kind, authoritative=True)
        except ValueError as exc:
            errors.append(str(exc))
            return []

    spf_records = []
    for raw in query(domain, 'TXT'):
        try:
            text = parse_txt(raw)
            if text.lower().startswith('v=spf1'):
                spf_records.append(text)
        except ValueError as exc:
            errors.append(str(exc))
    if len(spf_records) != 1 or spf_records[0].lower().split() != ['v=spf1', 'mx', '-all']:
        errors.append('SPF is not the approved single-host profile: v=spf1 mx -all')

    mx = query(domain, 'MX')
    if len(mx) != 1:
        errors.append(f'Expected exactly one MX pointing to {hostname}; found {mx}')
    else:
        parts = mx[0].split()
        if len(parts) != 2 or not parts[0].isdigit() or parts[1].rstrip('.').lower() != hostname:
            errors.append(f'MX is not the expected mailserver {hostname}: {mx[0]}')

    for kind, configured in [('A', bind4), ('AAAA', bind6)]:
        answers = query(hostname, kind)
        normalized = []
        for entry in answers:
            try:
                normalized.append(str(ipaddress.ip_address(entry)))
            except ValueError:
                errors.append(f'Invalid {kind} address returned for {hostname}: {entry}')
        if not configured:
            errors.append(f'Postfix smtp_bind_address{("6" if kind == "AAAA" else "")} is empty; cannot verify outbound IP')
        elif configured not in normalized:
            errors.append(f'Postfix outbound {kind} address {configured} is not published for {hostname}: {normalized}')

    print(f'   Authoritative NS {ns_name} ({ns_ip}): ' + ('OK' if not errors else 'FAILED'))
    for err in errors:
        print(f'      [ error ] {err}')
    return not errors


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--domain', required=True)
    ap.add_argument('--hostname', required=True)
    ap.add_argument('--ipv4', required=True)
    ap.add_argument('--ipv6', required=True)
    args = ap.parse_args()
    domain = args.domain.lower().rstrip('.')
    hostname = args.hostname.lower().rstrip('.')
    if not shutil.which('dig'):
        print('   [ error ] dig is required to verify public DNS')
        return 1
    try:
        bind4 = str(ipaddress.ip_address(args.ipv4)) if args.ipv4 else ''
        bind6 = str(ipaddress.ip_address(args.ipv6)) if args.ipv6 else ''
        zone, nameservers = discover_nameservers(domain)
        print(f'   DNS delegation zone................: {zone}')
        print(f'   Expected MX host...................: {hostname}')
        print(f'   Expected SMTP IPv4.................: {bind4 or "unknown"}')
        print(f'   Expected SMTP IPv6.................: {bind6 or "unknown"}')
        if len(nameservers) > 8:
            raise ValueError('More than eight authoritative nameservers; manual review required')
        ok = True
        for ns in nameservers:
            try:
                ip = addr_of_ns(ns)
                if not check_server(ns, ip, domain, hostname, bind4, bind6):
                    ok = False
            except ValueError as exc:
                print(f'   Authoritative NS {ns}: FAILED')
                print(f'      [ error ] {exc}')
                ok = False
        print('   Public SRS DNS preflight...........: ' + ('OK' if ok else 'FAILED'))
        return 0 if ok else 1
    except ValueError as exc:
        print(f'   [ error ] {exc}')
        print('   Public SRS DNS preflight...........: FAILED')
        return 1


if __name__ == '__main__':
    sys.exit(main())