zonemaster-tools/zmtools/dnsutil.py

"""DNS plumbing shared by every tool: name validation, single queries against a
given server, an iterative resolver that walks down from the root, and the
public-address check that keeps the service from being pointed at the jail
network."""

import base64
import ipaddress
import random
import re
import socket
import time

import httpx
import dns.edns
import dns.exception
import dns.flags
import dns.message
import dns.name
import dns.query
import dns.rcode
import dns.rdatatype
import dns.resolver
import dns.reversename

TIMEOUT = 3.0
UA = 'zonemaster.es-tools/1.0 (+https://zonemaster.es)'

ROOT_SERVERS = [
    ('a.root-servers.net', '198.41.0.4'),
    ('b.root-servers.net', '170.247.170.2'),
    ('c.root-servers.net', '192.33.4.12'),
    ('d.root-servers.net', '199.7.91.13'),
    ('e.root-servers.net', '192.203.230.10'),
    ('f.root-servers.net', '192.5.5.241'),
    ('i.root-servers.net', '192.36.148.17'),
    ('k.root-servers.net', '193.0.14.129'),
    ('l.root-servers.net', '199.7.83.42'),
    ('m.root-servers.net', '202.12.27.33'),
]

_LABEL = re.compile(r'^(?!-)[a-z0-9_-]{1,63}(?<!-)$')


class BadInput(ValueError):
    pass


def normalize_domain(raw, allow_single_label=False):
    """Return the ASCII (punycode) form of a user-typed name, or raise BadInput."""
    if not raw:
        raise BadInput('empty name')
    name = raw.strip().lower()
    name = re.sub(r'^[a-z]+://', '', name).split('/')[0].split('?')[0]
    name = name.rstrip('.')
    if ':' in name:  # host:port
        name = name.split(':')[0]
    try:
        name = name.encode('idna').decode('ascii')
    except UnicodeError:
        raise BadInput('invalid internationalized name')
    labels = name.split('.')
    if len(name) > 253 or not labels or (len(labels) < 2 and not allow_single_label):
        raise BadInput('not a domain name')
    if not all(_LABEL.match(l) for l in labels):
        raise BadInput('not a domain name')
    if labels[-1].isdigit():
        raise BadInput('not a domain name')
    return name


def is_public_ip(addr):
    try:
        ip = ipaddress.ip_address(addr)
    except ValueError:
        return False
    return ip.is_global and not ip.is_multicast


def public_ip(raw):
    """Validate a user-supplied resolver address."""
    raw = (raw or '').strip()
    if not is_public_ip(raw):
        raise BadInput('resolver must be a public IP address')
    return str(ipaddress.ip_address(raw))


def make_query(name, rdtype, dnssec=False, ecs=None, rd=True):
    options = []
    if ecs:
        addr, plen = ecs
        options.append(dns.edns.ECSOption(addr, plen))
    q = dns.message.make_query(name, rdtype, want_dnssec=dnssec, use_edns=0,
                               payload=1232, options=options)
    if not rd:
        q.flags &= ~dns.flags.RD
    return q


def query(server, name, rdtype, *, tcp=False, dnssec=False, ecs=None, rd=True,
          cd=False, timeout=TIMEOUT):
    """One query to one server. Returns (response|None, elapsed_ms, error|None).
    Falls back to TCP on truncation."""
    q = make_query(name, rdtype, dnssec=dnssec, ecs=ecs, rd=rd)
    if cd:
        q.flags |= dns.flags.CD
    t0 = time.monotonic()
    try:
        if tcp:
            r = dns.query.tcp(q, server, timeout=timeout)
        else:
            r = dns.query.udp(q, server, timeout=timeout, ignore_unexpected=True)
            if r.flags & dns.flags.TC:
                r = dns.query.tcp(q, server, timeout=timeout)
        return r, round((time.monotonic() - t0) * 1000, 1), None
    except dns.exception.Timeout:
        return None, round((time.monotonic() - t0) * 1000, 1), 'timeout'
    except (OSError, dns.exception.DNSException) as e:
        return None, round((time.monotonic() - t0) * 1000, 1), type(e).__name__


# Several DoH services (Quad9, CIRA, Mullvad, CZ.NIC) require HTTP/2, as
# RFC 8484 recommends, so this needs httpx + h2 rather than urllib.
_doh = httpx.Client(http2=True, timeout=TIMEOUT + 2, follow_redirects=False,
                    headers={'Accept': 'application/dns-message', 'User-Agent': UA},
                    limits=httpx.Limits(max_connections=64, max_keepalive_connections=16))


def query_doh(url, name, rdtype, dnssec=True):
    """RFC 8484 GET. Returns (response|None, elapsed_ms, error|None)."""
    q = make_query(name, rdtype, dnssec=dnssec)
    q.id = 0
    b64 = base64.urlsafe_b64encode(q.to_wire()).decode().rstrip('=')
    t0 = time.monotonic()
    try:
        resp = _doh.get(url, params={'dns': b64})
        resp.raise_for_status()
        return dns.message.from_wire(resp.content), round((time.monotonic() - t0) * 1000, 1), None
    except Exception as e:  # network, HTTP or parse errors all mean "no answer"
        return None, round((time.monotonic() - t0) * 1000, 1), type(e).__name__


def answers(resp, rdtype=None):
    """Answer-section records as [{'name','type','ttl','value'}]."""
    out = []
    if resp is None:
        return out
    for rrset in resp.answer:
        if rdtype is not None and rrset.rdtype not in (rdtype, dns.rdatatype.CNAME):
            continue
        if rrset.rdtype == dns.rdatatype.RRSIG:
            continue
        for rd in rrset:
            out.append({'name': rrset.name.to_text(omit_final_dot=True),
                        'type': dns.rdatatype.to_text(rrset.rdtype),
                        'ttl': rrset.ttl, 'value': rd.to_text()})
    return out


def rcode_text(resp):
    return dns.rcode.to_text(resp.rcode()) if resp is not None else None


# A recursive resolver for the tools' own lookups (not the thing under test).
_stub = dns.resolver.Resolver(configure=False)
_stub.nameservers = ['1.1.1.1', '9.9.9.9', '8.8.8.8']
_stub.timeout = TIMEOUT
_stub.lifetime = TIMEOUT * 2
_stub.use_edns(0, 0, 1232)


def resolve(name, rdtype):
    """Values of an RRset via the stub, [] on NXDOMAIN/NODATA/failure."""
    try:
        ans = _stub.resolve(name, rdtype, raise_on_no_answer=False)
    except (dns.resolver.NXDOMAIN, dns.resolver.NoNameservers,
            dns.exception.Timeout, dns.resolver.NoAnswer):
        return []
    if ans.rrset is None:
        return []
    return [rd.to_text() for rd in ans.rrset]


def resolve_full(name, rdtype):
    """(rcode_text, [values], ttl) via the stub."""
    try:
        ans = _stub.resolve(name, rdtype, raise_on_no_answer=False)
    except dns.resolver.NXDOMAIN:
        return 'NXDOMAIN', [], None
    except (dns.resolver.NoNameservers, dns.exception.Timeout):
        return 'SERVFAIL', [], None
    if ans.rrset is None:
        return 'NOERROR', [], None
    return 'NOERROR', [rd.to_text() for rd in ans.rrset], ans.rrset.ttl


def txt_values(name):
    vals = []
    for t in resolve(name, 'TXT'):
        # '"part1" "part2"' -> 'part1part2'
        vals.append(''.join(re.findall(r'"((?:[^"\\]|\\.)*)"', t)) or t)
    return vals


def addresses(name, v6=True):
    out = resolve(name, 'A')
    if v6:
        out += resolve(name, 'AAAA')
    return out


def cname_chain(name, limit=8):
    chain, cur = [], name
    for _ in range(limit):
        nxt = resolve(cur, 'CNAME')
        if not nxt:
            break
        cur = nxt[0].rstrip('.')
        chain.append(cur)
    return chain


def _referral(resp):
    """From a non-authoritative response: (zone, [ns names], {ns: [glue ips]})."""
    zone, names = None, []
    for rrset in resp.authority:
        if rrset.rdtype == dns.rdatatype.NS:
            zone = rrset.name.to_text(omit_final_dot=True)
            names = [rd.target.to_text(omit_final_dot=True) for rd in rrset]
    glue = {}
    for rrset in resp.additional:
        if rrset.rdtype == dns.rdatatype.A:
            glue.setdefault(rrset.name.to_text(omit_final_dot=True), []).extend(
                rd.address for rd in rrset)
    return zone, names, glue


def delegation_chain(domain, rdtype='NS', max_steps=12):
    """Walk from the root to the zone holding `domain`, like `dig +trace`.

    Returns a list of steps: {'zone', 'server', 'server_ip', 'ns': [...],
    'glue': {...}, 'ds': bool, 'ms'} — one per zone cut — and the final
    authoritative step with 'authoritative': True."""
    steps = []
    servers = random.sample(ROOT_SERVERS, 3)
    zone = '.'
    for _ in range(max_steps):
        resp = None
        for host, ip in servers:
            resp, ms, err = query(ip, domain, rdtype, rd=False, dnssec=True)
            if resp is not None:
                break
        if resp is None:
            steps.append({'zone': zone, 'error': 'no server answered'})
            return steps
        auth = bool(resp.flags & dns.flags.AA) or (
            resp.answer and not any(r.rdtype == dns.rdatatype.NS for r in resp.authority))
        if auth or resp.rcode() != dns.rcode.NOERROR:
            steps.append({'zone': zone, 'server': host, 'server_ip': ip, 'ms': ms,
                          'authoritative': True, 'rcode': rcode_text(resp),
                          'answer': answers(resp)})
            return steps
        nzone, names, glue = _referral(resp)
        if not names or nzone is None or nzone == zone.rstrip('.'):
            steps.append({'zone': zone, 'server': host, 'server_ip': ip, 'ms': ms,
                          'error': 'lame or empty referral'})
            return steps
        has_ds = any(r.rdtype == dns.rdatatype.DS for r in resp.authority)
        steps.append({'zone': zone, 'server': host, 'server_ip': ip, 'ms': ms,
                      'child': nzone, 'ns': names, 'glue': glue, 'ds': has_ds})
        nxt = [(n, glue[n][0]) for n in names if glue.get(n)]
        if not nxt:  # out-of-bailiwick NS without glue: resolve them
            for n in names[:4]:
                ips = resolve(n, 'A')
                if ips:
                    nxt.append((n, ips[0]))
        nxt = [(n, i) for n, i in nxt if is_public_ip(i)]
        if not nxt:
            steps.append({'zone': nzone, 'error': 'nameservers do not resolve'})
            return steps
        random.shuffle(nxt)
        servers = nxt[:3]
        zone = nzone
    return steps


_auth_cache = {}


def zone_servers(zone):
    """Public IPv4 addresses of a zone's nameservers (cached 1 h)."""
    hit = _auth_cache.get(zone)
    if hit and hit[0] > time.time():
        return hit[1]
    ips, name = [], zone
    while not ips and name.count('.') >= 1:  # not delegated? ask the parent
        for ns in resolve(name, 'NS'):
            ips += [i for i in resolve(ns.rstrip('.'), 'A') if is_public_ip(i)]
        name = name.split('.', 1)[1]
    _auth_cache[zone] = (time.time() + 3600, ips)
    return ips


def query_authoritative(zone, name, rdtype='A', hops=4):
    """Ask the zone's own nameservers directly, following referrals to any
    deeper delegation — used for DNSBLs, which refuse or poison answers sent
    through big public resolvers. Returns (rcode_text|None, [values])."""
    ips = list(zone_servers(zone))
    for _ in range(hops):
        random.shuffle(ips)
        resp = None
        for ip in ips[:3]:
            resp, _, err = query(ip, name, rdtype, rd=False, timeout=2.5)
            if resp is not None:
                break
        if resp is None:
            return None, []
        vals = [a['value'] for a in answers(resp, dns.rdatatype.from_text(rdtype))
                if a['type'] == rdtype]
        if vals or resp.flags & dns.flags.AA or resp.rcode() != dns.rcode.NOERROR:
            return rcode_text(resp), vals
        _, names, glue = _referral(resp)
        if not names:
            return rcode_text(resp), vals
        ips = [i for n in names for i in (glue.get(n) or resolve(n, 'A')) if is_public_ip(i)]
        if not ips:
            return None, []
    return None, []


def reverse_name(ip):
    return dns.reversename.from_address(ip).to_text(omit_final_dot=True)


def ptr(ip):
    vals = resolve(reverse_name(ip), 'PTR')
    return vals[0].rstrip('.') if vals else None


def socket_family(ip):
    return socket.AF_INET6 if ':' in ip else socket.AF_INET