netbox-plugin-dns-bridge/netbox_dns_bridge/tests/_helpers.py

"""Shared, non-fixture helpers for the test suite.

Pure DNS / TSIG / handler plumbing that does not touch the database. Database
factories live in ``conftest.py`` as pytest fixtures.
"""
from __future__ import annotations

from unittest.mock import MagicMock

import dns.message
import dns.name
import dns.rdatatype
import dns.tsig
import dns.tsigkeyring

from netbox_dns_bridge.management.commands.dns_transfer_endpoint.request_handler import (
    DNSBaseRequestHandler,
    TCPRequestHandler,
    UDPRequestHandler,
)

DEFAULT_TSIG_ALGORITHM = "hmac-sha256"
DEFAULT_TSIG_SECRET = "YWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWFhYWE="


def make_tsig_key(
    key_name: str = "testkey.",
    secret: str = DEFAULT_TSIG_SECRET,
    algorithm: str = DEFAULT_TSIG_ALGORITHM,
) -> dns.tsig.Key:
    name = dns.name.from_text(key_name, origin=dns.name.root).canonicalize()
    return dns.tsig.Key(name=name, secret=secret, algorithm=algorithm)


def build_keyring(tsig_key: dns.tsig.Key) -> dict:
    return {tsig_key.name: tsig_key}


def build_tsig_view_map(key_name: str, view) -> dict:
    canon = dns.name.from_text(key_name, origin=dns.name.root).canonicalize().to_text()
    return {canon: view}


def make_unsigned_query(zone_name: str, rdtype) -> bytes:
    qname = dns.name.from_text(zone_name, dns.name.root)
    msg = dns.message.make_query(qname, rdtype)
    return msg.to_wire()


def make_signed_query(
    zone_name: str,
    rdtype,
    tsig_key: dns.tsig.Key,
) -> bytes:
    qname = dns.name.from_text(zone_name, dns.name.root)
    msg = dns.message.make_query(qname, rdtype)
    msg.use_tsig(keyring={tsig_key.name: tsig_key}, keyname=tsig_key.name)
    return msg.to_wire()


class FakeServer:
    """Stand-in for socketserver's ``self.server`` carrying TSIG state."""

    def __init__(self, keyring: dict, tsig_view_map: dict):
        self.keyring = keyring
        self.tsig_view_map = tsig_view_map


class _CaptureMixin:
    """Capture every wire message a handler emits.

    Both ``_send_response`` (used by denials, SOA and the final AXFR message)
    and the direct ``self.request.sendall`` calls (intermediate AXFR messages)
    are recorded as raw wire bytes (the 2-byte TCP length prefix is stripped).
    """

    transport = "tcp"

    def _setup_capture(self, captured: list):
        self._captured = captured
        # For TCP: AXFR intermediate sends use self.request.sendall directly.
        self.request = MagicMock()
        self.request.sendall.side_effect = lambda data: self._record_tcp_sendall(data)

    def _record_tcp_sendall(self, data: bytes):
        # Intermediate AXFR sends are ``len(wire).to_bytes(2,"big") + wire``.
        if len(data) >= 2:
            self._captured.append(data[2:])
        else:
            self._captured.append(data)

    def _send_response(self, data: bytes):
        self._captured.append(data)


class CaptureTCPHandler(_CaptureMixin, TCPRequestHandler):
    transport = "tcp"

    def handle(self):
        # Never invoked: tests call _handle_dns_query directly.
        return


class CaptureUDPHandler(_CaptureMixin, UDPRequestHandler):
    transport = "udp"

    def handle(self):
        return


def make_capture_handler(
    handler_cls,
    keyring: dict,
    tsig_view_map: dict,
    client_ip: str = "127.0.0.1",
):
    """Instantiate a capture handler without running ``handle()``.

    Returns ``(handler, captured)`` where ``captured`` is the list of wire
    messages emitted by the handler.
    """
    captured: list = []
    server = FakeServer(keyring, tsig_view_map)
    handler = handler_cls.__new__(handler_cls)
    # Replicate DNSBaseRequestHandler.__init__ without calling super().__init__
    # (which would invoke setup()/handle()).
    handler.MAX_WIRE = 65535
    handler.RESERVED_TSIG = 300
    handler._setup_capture(captured)
    handler.client_address = (client_ip, 12345)
    handler.server = server
    return handler, captured


def parse_response(wire: bytes, keyring=None, one_rr_per_rrset: bool = False) -> dns.message.Message:
    # ``make_response`` in this dnspython copies the request TSIG keyname into
    # the response even for unsigned denials, so lenient parsing is required to
    # inspect rcode/flags/answer without verifying the (absent) MAC.
    # Use ``one_rr_per_rrset=True`` for AXFR framing checks so that repeated
    # identical RRsets (e.g. opening + closing SOA) are not merged on parse.
    return dns.message.from_wire(
        wire,
        keyring=keyring or {},
        continue_on_error=True,
        ignore_trailing=True,
        one_rr_per_rrset=one_rr_per_rrset,
    )


def parse_response_strict(
    wire: bytes,
    keyring: dict,
    request_mac=None,
    one_rr_per_rrset: bool = False,
) -> dns.message.Message:
    """Parse and fully verify TSIG. Raises on bad signature / unknown key.

    Verifying a response TSIG requires the original request's MAC; pass it via
    ``request_mac`` (obtainable from the parsed query's ``.mac`` attribute).
    """
    return dns.message.from_wire(
        wire,
        keyring=keyring,
        request_mac=request_mac,
        ignore_trailing=True,
        one_rr_per_rrset=one_rr_per_rrset,
    )


def query_mac(query_wire: bytes, keyring: dict) -> bytes:
    """Return the request MAC of a signed query wire (for response verification)."""
    q = dns.message.from_wire(query_wire, keyring=keyring, ignore_trailing=True)
    return q.mac


def tcp_exchange(wire: bytes, host: str, port: int, timeout: float = 5.0) -> bytes:
    """Send a wire message over TCP (2-byte length framing), return response wire."""
    import socket

    with socket.create_connection((host, port), timeout=timeout) as sock:
        sock.settimeout(timeout)
        sock.sendall(len(wire).to_bytes(2, "big") + wire)
        # Read response length then body (handle partial reads).
        length = _recv_exact(sock, 2)
        resp_len = int.from_bytes(length, "big")
        return _recv_exact(sock, resp_len)


def udp_exchange(wire: bytes, host: str, port: int, timeout: float = 5.0) -> bytes:
    """Send a wire message over UDP, return response wire."""
    import socket

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(timeout)
    try:
        sock.sendto(wire, (host, port))
        return sock.recv(65535)
    finally:
        sock.close()


def _recv_exact(sock, n: int) -> bytes:
    buf = b""
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk:
            raise OSError("connection closed")
        buf += chunk
    return buf


def xfr_rrsets(xfr_messages):
    """Flatten AXFR/IXFR response messages into a flat list of RRsets.

    ``dns.query.xfr`` yields ``Message`` objects (one per TCP response chunk).
    Each message's ``answer`` section contains RRsets. This helper flattens
    all answer RRsets across all messages into a single list, preserving order.
    """
    rrsets = []
    for msg in xfr_messages:
        rrsets.extend(msg.answer)
    return rrsets


def make_dns_zone(name: str, soa_serial: int = 1, soa_ttl: int = 300,
                  extra=None) -> "dns.zone.Zone":
    """Build a ``dns.zone.Zone`` directly (no database) for unit tests.

    Always includes an SOA at the apex. ``extra`` is an optional list of
    ``(relative_name, rdtype, ttl, value_text)`` tuples appended as rdatasets.
    """
    import dns.rdata
    import dns.rdataclass
    import dns.rdataset
    import dns.rdatatype
    import dns.zone

    origin = dns.name.from_text(name, dns.name.root)
    zone = dns.zone.Zone(origin)
    zone.rdclass = dns.rdataclass.IN

    soa_rdata = dns.rdata.from_text(
        dns.rdataclass.IN,
        dns.rdatatype.SOA,
        f"ns1.{name}. hostmaster.{name}. {soa_serial} 300 100 1209600 3600",
    )
    soa_rds = dns.rdataset.Rdataset(dns.rdataclass.IN, dns.rdatatype.SOA)
    soa_rds.add(soa_rdata, soa_ttl)
    zone.replace_rdataset(origin, soa_rds)

    for rel_name, rdtype, ttl, value_text in extra or []:
        owner = dns.name.from_text(rel_name, origin)
        rdtype_int = dns.rdatatype.from_text(rdtype)
        rdata = dns.rdata.from_text(dns.rdataclass.IN, rdtype_int, value_text)
        rds = zone.find_rdataset(owner, rdtype_int, create=True)
        rds.add(rdata, ttl)

    return zone