"""DNS message (RFC 1035, plus common extensions).
Domain-name compression is the one thing here that doesn't fit rustruct's
declarative model: a pointer's target is an absolute offset into the whole
message, and writing one requires remembering every name suffix already
emitted anywhere earlier in the message -- state that outlives any single
field or nested scope. So the record types that carry a domain name
(`Question`, `ResourceRecord`, and the NS/CNAME/DNAME/SOA/PTR/MX/SRV/NAPTR/
RRSIG/NSEC/RP/HTTPS record types) are small hand-written classes operating
on a shared `bytearray` + a `dict` context, calling `write_name`/
`decode_name` for each name and `encode_labels` to prepare one for the
wire.
Everything else about those same record types -- every fixed-width field
or run of them, both before and after the domain name(s) -- is still a
real, nested `rustruct.Struct` (`SOAFixed`, `SRVFixed`, `RRSIGFixed`, and
so on): there is no hand-rolled `struct.pack()` format string anywhere in
this module, right up through the 12-byte message header itself
(`MessageHeader`, nesting `DNSFlags`). Record types with no domain name at
all (`A`/`AAAA`/`DS`/`DNSKEY`/`TLSA`/`SSHFP`/`CAA`/`LOC`) need no escape
hatch whatsoever and are plain top-level `Struct` classes.
Domain names are plain strings (`"example.com"`; the root is `""`; a
trailing dot is accepted and dropped). Labels are transcoded byte-for-byte
as latin-1. Compression pointers are followed on read (backwards only, so
hostile messages cannot loop) and emitted on write whenever a name suffix
was already written; pack with `compress=False` to disable that.
A handful of record types carry a domain name that RFC 4034/2782 forbid
compressing (RRSIG's signer, NSEC's next-domain, SRV's/HTTPS's target):
those call `write_name(out, labels, None)` regardless of the message's own
`compress=` setting, matching what real resolvers expect on the wire.
"""
import enum
from dataclasses import dataclass, field
from typing import Any, Protocol
from rustruct import U8, U16, U32, Struct, bits, convert, described, slice, string
from .inet import ipv4_address_field, ipv6_address_field
__all__ = [
"OpenIntEnum",
"RRType",
"RCODE",
"DNSClass",
"DNSFlags",
"A",
"AAAA",
"NS",
"CNAME",
"DNAME",
"PTR",
"MX",
"SOA",
"SRV",
"NAPTR",
"DS",
"DNSKEY",
"RRSIG",
"NSEC",
"CAA",
"HTTPS",
"LOC",
"RP",
"TLSA",
"SSHFP",
"TXT",
"RDataKind",
"UnknownRData",
"Question",
"ResourceRecord",
"DNS",
"reply",
"edns0",
"edns_udp_size",
]
MAX_NAME_WIRE = 255 # labels + length bytes + the final zero (RFC 1035 2.3.4)
MAX_LABEL = 63
POINTER_LIMIT = 0x3FFF # compression offsets are 14 bits
class OpenIntEnum(enum.IntEnum):
"""An `IntEnum` that tolerates values with no named member: DNS wire
rtype/qtype/qclass values aren't restricted to well-known ones, so
decoding one must not raise -- it behaves as a plain int with no name.
Unknown values are cached on `PSEUDO` so repeated decodes of the same
value return the identical object."""
@classmethod
def _missing_(cls, value: object) -> "OpenIntEnum | None":
if not isinstance(value, int) or not 0 <= value < 0x10000:
return None
cache = cls.__dict__.get("PSEUDO")
if cache is None:
cache = cls.PSEUDO = {}
member = cache.get(value)
if member is None:
member = int.__new__(cls, value)
member._name_ = f"{cls.__name__}_{value}"
member._value_ = value
cache[value] = member
return member
[docs]
class RRType(OpenIntEnum):
A = 1
NS = 2
CNAME = 5
SOA = 6
PTR = 12
RP = 17
MX = 15
TXT = 16
AAAA = 28
LOC = 29
SRV = 33
NAPTR = 35
DNAME = 39
OPT = 41
DS = 43
SSHFP = 44
RRSIG = 46
NSEC = 47
DNSKEY = 48
TLSA = 52
HTTPS = 65
CAA = 257
ANY = 255
[docs]
class DNSClass(OpenIntEnum):
IN = 1
CH = 3
HS = 4
ANY = 255
class RCODE(OpenIntEnum):
"""RFC 1035 section 4.1.1 plus the extensions RFC 2136/2671 added; the
header's own `rcode` field is only 4 bits wide (0-15), the rest need an
EDNS0 extended-RCODE to appear on the wire at all."""
NOERROR = 0
FORMERR = 1
SERVFAIL = 2
NXDOMAIN = 3
NOTIMP = 4
REFUSED = 5
YXDOMAIN = 6
YXRRSET = 7
NXRRSET = 8
NOTAUTH = 9
NOTZONE = 10
[docs]
class DNSFlags(Struct, byteorder="big"):
"""The second header word: QR, opcode, AA/TC/RD/RA, Z and RCODE."""
qr: bool = bits(1, default=False, help="query (False) or response (True)")
opcode: int = bits(4, default=0, help="kind of query; 0 = standard query")
aa: bool = bits(1, default=False, help="responder is authoritative")
tc: bool = bits(1, default=False, help="message was truncated")
rd: bool = bits(1, default=False, help="recursion desired")
ra: bool = bits(1, default=False, help="recursion available")
z: int = bits(3, default=0, help="reserved; must be zero")
rcode: int = bits(4, default=0, help="response code; 0 = no error")
[docs]
class A(Struct, byteorder="big"):
"""No domain name -- a plain 4-byte address, fully declarative."""
RTYPE = RRType.A
address: object = ipv4_address_field(help="the host's IPv4 address")
[docs]
class AAAA(Struct, byteorder="big"):
RTYPE = RRType.AAAA
address: object = ipv6_address_field(help="the host's IPv6 address")
class DS(Struct, byteorder="big"):
"""Delegation Signer (RFC 4034 section 5); no domain name involved."""
RTYPE = RRType.DS
key_tag: U16
algorithm: U8
digest_type: U8
digest: bytes = slice(len="*")
class DNSKEY(Struct, byteorder="big"):
"""RFC 4034 section 2; no domain name involved."""
RTYPE = RRType.DNSKEY
flags: U16
protocol: U8
algorithm: U8
key: bytes = slice(len="*")
class TLSA(Struct, byteorder="big"):
"""RFC 6698; no domain name involved."""
RTYPE = RRType.TLSA
usage: U8
selector: U8
matching_type: U8
cert_data: bytes = slice(len="*")
class SSHFP(Struct, byteorder="big"):
"""RFC 4255; no domain name involved."""
RTYPE = RRType.SSHFP
algorithm: U8
fp_type: U8
fingerprint: bytes = slice(len="*")
class CAA(Struct, byteorder="big"):
"""RFC 6844; no domain name involved. `tag_length` is derived from
`len(tag)` on pack, like a length field -- the caller never supplies it
(see `described()`'s docstring)."""
RTYPE = RRType.CAA
flags: U8 = described(default=0, help="critical flag; bit 0 set means non-issuer-critical")
tag_length: U8 = described(help="derived from len(tag) on pack")
tag: str = string(len="tag_length", encoding="ascii")
value: str = string(len="*", encoding="utf-8")
def loc_precision_to_wire(value: float) -> int:
"""RFC 1876's SIZE/HORIZ_PRE/VERT_PRE nibble-exponent encoding: value ==
0 packs as the single byte 0, otherwise `(mantissa << 4) | exponent`
with `mantissa * 10**exponent == round(value * 100)` and mantissa a
single digit 0-9."""
if value == 0:
return 0
exponent = 0
scaled = value * 100
while scaled >= 10 and exponent < 9:
scaled /= 10
exponent += 1
mantissa = int(round(scaled))
if mantissa >= 10:
raise ValueError(f"LOC precision value out of range: {value!r}")
return (mantissa << 4) | exponent
def loc_precision_from_wire(value: int) -> float:
mantissa, exponent = value >> 4, value & 0x0F
return mantissa * (10**exponent) / 100
def loc_coord_to_wire(degrees: float) -> int:
return int(round(degrees * 3600000)) + (1 << 31)
def loc_coord_from_wire(value: int) -> float:
return (value - (1 << 31)) / 3600000
def loc_altitude_to_wire(meters: float) -> int:
return int(round((meters + 100000) * 100))
def loc_altitude_from_wire(value: int) -> float:
return value / 100 - 100000
def loc_precision_field(**kwargs: Any) -> Any:
return convert(U8, decode=loc_precision_from_wire, encode=loc_precision_to_wire, **kwargs)
class LOC(Struct, byteorder="big"):
"""RFC 1876; no domain name involved. Latitude/longitude are degrees
(negative is south/west), altitude is meters above the WGS 84 reference
ellipsoid, and size/*_precision are meters -- all four wire fixed-point
encodings (see the `loc_*` helpers above) round-trip through plain
Python floats via `convert()`."""
RTYPE = RRType.LOC
version: U8 = described(default=0, help="always 0 (RFC 1876 has no other version)")
size: float = loc_precision_field(default=1.0, help="diameter of the enclosing sphere, in meters")
h_precision: float = loc_precision_field(default=10000.0, help="horizontal precision, in meters")
v_precision: float = loc_precision_field(default=10.0, help="vertical precision, in meters")
latitude: float = convert(U32, decode=loc_coord_from_wire, encode=loc_coord_to_wire, help="degrees north")
longitude: float = convert(U32, decode=loc_coord_from_wire, encode=loc_coord_to_wire, help="degrees east")
altitude: float = convert(U32, decode=loc_altitude_from_wire, encode=loc_altitude_to_wire)
# ---------- domain names: the hand-written escape hatch ----------
def pointer_target(buf: bytes | bytearray, off: int) -> int:
if off + 1 >= len(buf):
raise ValueError(f"cut compression pointer at offset {off}")
target = ((buf[off] & 0x3F) << 8) | buf[off + 1]
if target >= off:
raise ValueError(f"compression pointer at offset {off} does not point backwards")
return target
def read_label(buf: bytes | bytearray, off: int, length: int) -> tuple[str, int]:
stop = off + 1 + length
label = bytes(buf[off + 1 : stop])
if len(label) != length:
raise ValueError(f"cut label at offset {off}")
return label.decode("latin-1"), stop
def decode_name(buf: bytes | bytearray, start: int) -> tuple[str, int]:
"""Returns (name, offset to resume the *enclosing* record at) -- the
resume point is right after the first pointer followed, not wherever the
pointer chain eventually bottoms out."""
off = start
resume = None
labels = []
wire_len = 1 # the terminating zero
while True:
if off >= len(buf):
raise ValueError(f"domain name runs past the buffer at offset {off}")
first = buf[off]
if first == 0:
off += 1
break
kind = first & 0xC0
if kind == 0xC0:
target = pointer_target(buf, off)
if resume is None:
resume = off + 2
off = target
continue
if kind:
raise ValueError(f"unsupported label type 0x{first:02X} at offset {off}")
label, off = read_label(buf, off, first)
wire_len += first + 1
if wire_len > MAX_NAME_WIRE:
raise ValueError("domain name longer than 255 wire bytes")
labels.append(label)
return ".".join(labels), (resume if resume is not None else off)
def encode_labels(value: str) -> list[bytes]:
stripped = value[:-1] if value.endswith(".") else value
labels = [part.encode("latin-1") for part in stripped.split(".")] if stripped else []
wire_len = 1
for label in labels:
if not 1 <= len(label) <= MAX_LABEL:
raise ValueError(f"bad label length {len(label)} in {value!r}")
wire_len += len(label) + 1
if wire_len > MAX_NAME_WIRE:
raise ValueError(f"domain name {value!r} longer than 255 wire bytes")
return labels
def write_name(out: bytearray, labels: list[bytes], ctx: dict[str, Any] | None) -> None:
"""`ctx`, if not None, is the whole message's shared compression table
(`{"dns_names": {label_suffix_tuple: offset}}`); pass None to disable
compression on write (`DNS.pack(compress=False)`)."""
offsets = ctx.setdefault("dns_names", {}) if ctx is not None else None
for i, label in enumerate(labels):
if offsets is not None:
suffix = tuple(labels[i:])
target = offsets.get(suffix)
if target is not None:
out += bytes((0xC0 | (target >> 8), target & 0xFF))
return
if len(out) <= POINTER_LIMIT:
offsets[suffix] = len(out)
out.append(len(label))
out += label
out.append(0)
# ---------- RDATA types that contain a domain name: hand-written ----------
@dataclass
class SingleName:
"""Common shape shared by NS/CNAME/PTR/DNAME (RFC 1035/6672): a
compressible domain name and nothing else. Concrete subclasses add
only their own `RTYPE`; `pack_into`/`unpack` and equality/repr are
inherited as-is, with `cls(...)` in `unpack` dispatching to whichever
subclass it's called on."""
target: str
def pack_into(self, out: bytearray, ctx: dict[str, Any] | None) -> None:
write_name(out, encode_labels(self.target), ctx)
@classmethod
def unpack(cls, buf: bytes | bytearray, off: int, length: int, ctx: dict[str, Any]) -> "SingleName":
name, _ = decode_name(buf, off)
return cls(name)
[docs]
class NS(SingleName):
RTYPE = RRType.NS
[docs]
class CNAME(SingleName):
RTYPE = RRType.CNAME
[docs]
class PTR(SingleName):
RTYPE = RRType.PTR
class DNAME(SingleName):
"""RFC 6672; a name-compressible domain name, same as CNAME/NS/PTR."""
RTYPE = RRType.DNAME
class MXPrefix(Struct, byteorder="big"):
"""MX's fixed-width part, ahead of the `exchange` domain name."""
preference: U16
[docs]
@dataclass
class MX:
RTYPE = RRType.MX
preference: int
exchange: str
def pack_into(self, out: bytearray, ctx: dict[str, Any] | None) -> None:
out += MXPrefix(preference=self.preference).pack()
write_name(out, encode_labels(self.exchange), ctx)
@classmethod
def unpack(cls, buf: bytes | bytearray, off: int, length: int, ctx: dict[str, Any]) -> "MX":
prefix, pos = MXPrefix.unpack_from(buf, off)
exchange, _ = decode_name(buf, pos)
return cls(prefix.preference, exchange)
class SOAFixed(Struct, byteorder="big"):
"""SOA's five fixed-width fields, after `mname`/`rname`."""
serial: U32
refresh: U32
retry: U32
expire: U32
minimum: U32
[docs]
@dataclass
class SOA:
RTYPE = RRType.SOA
mname: str
rname: str
serial: int
refresh: int
retry: int
expire: int
minimum: int
def pack_into(self, out: bytearray, ctx: dict[str, Any] | None) -> None:
write_name(out, encode_labels(self.mname), ctx)
write_name(out, encode_labels(self.rname), ctx)
out += SOAFixed(
serial=self.serial, refresh=self.refresh, retry=self.retry, expire=self.expire, minimum=self.minimum
).pack()
@classmethod
def unpack(cls, buf: bytes | bytearray, off: int, length: int, ctx: dict[str, Any]) -> "SOA":
mname, off = decode_name(buf, off)
rname, off = decode_name(buf, off)
fixed, _ = SOAFixed.unpack_from(buf, off)
return cls(mname, rname, fixed.serial, fixed.refresh, fixed.retry, fixed.expire, fixed.minimum)
[docs]
@dataclass
class TXT:
"""A sequence of length-prefixed opaque text chunks filling the record."""
RTYPE = RRType.TXT
strings: list[bytes] = field(default_factory=list)
def pack_into(self, out: bytearray, ctx: dict[str, Any] | None) -> None:
for s in self.strings:
out.append(len(s))
out += s
@classmethod
def unpack(cls, buf: bytes | bytearray, off: int, length: int, ctx: dict[str, Any]) -> "TXT":
strings = []
end = off + length
pos = off
while pos < end:
n = buf[pos]
pos += 1
strings.append(bytes(buf[pos : pos + n]))
pos += n
return cls(strings)
class SRVFixed(Struct, byteorder="big"):
"""SRV's three fixed-width fields, ahead of `target`."""
priority: U16
weight: U16
port: U16
@dataclass
class SRV:
"""RFC 2782; `target` must not be compressed."""
RTYPE = RRType.SRV
priority: int
weight: int
port: int
target: str
def pack_into(self, out: bytearray, ctx: dict[str, Any] | None) -> None:
out += SRVFixed(priority=self.priority, weight=self.weight, port=self.port).pack()
write_name(out, encode_labels(self.target), None)
@classmethod
def unpack(cls, buf: bytes | bytearray, off: int, length: int, ctx: dict[str, Any]) -> "SRV":
fixed, pos = SRVFixed.unpack_from(buf, off)
target, _ = decode_name(buf, pos)
return cls(fixed.priority, fixed.weight, fixed.port, target)
def read_char_string(buf: bytes | bytearray, off: int) -> tuple[str, int]:
"""A DNS "character-string" (RFC 1035 section 3.3): one length octet
followed by that many bytes, decoded as ASCII. Returns (text, offset to
resume at)."""
n = buf[off]
stop = off + 1 + n
return bytes(buf[off + 1 : stop]).decode("ascii"), stop
def write_char_string(out: bytearray, value: str) -> None:
encoded = value.encode("ascii")
out.append(len(encoded))
out += encoded
class NAPTRPrefix(Struct, byteorder="big"):
"""NAPTR's two fixed-width fields, ahead of its character-strings."""
order: U16
preference: U16
@dataclass
class NAPTR:
"""RFC 3403; `replacement` is name-compressible, same as CNAME/NS/PTR."""
RTYPE = RRType.NAPTR
order: int
preference: int
flags: str
service: str
regexp: str
replacement: str
def pack_into(self, out: bytearray, ctx: dict[str, Any] | None) -> None:
out += NAPTRPrefix(order=self.order, preference=self.preference).pack()
write_char_string(out, self.flags)
write_char_string(out, self.service)
write_char_string(out, self.regexp)
write_name(out, encode_labels(self.replacement), ctx)
@classmethod
def unpack(cls, buf: bytes | bytearray, off: int, length: int, ctx: dict[str, Any]) -> "NAPTR":
prefix, pos = NAPTRPrefix.unpack_from(buf, off)
flags, pos = read_char_string(buf, pos)
service, pos = read_char_string(buf, pos)
regexp, pos = read_char_string(buf, pos)
replacement, _ = decode_name(buf, pos)
return cls(prefix.order, prefix.preference, flags, service, regexp, replacement)
class RRSIGFixed(Struct, byteorder="big"):
"""RRSIG's seven fixed-width fields, ahead of `signer`."""
type_covered: U16
algorithm: U8
labels: U8
original_ttl: U32
expiration: U32
inception: U32
key_tag: U16
@dataclass
class RRSIG:
"""RFC 4034 section 3; `signer` must not be compressed."""
RTYPE = RRType.RRSIG
type_covered: int
algorithm: int
labels: int
original_ttl: int
expiration: int
inception: int
key_tag: int
signer: str
signature: bytes
def pack_into(self, out: bytearray, ctx: dict[str, Any] | None) -> None:
out += RRSIGFixed(
type_covered=self.type_covered,
algorithm=self.algorithm,
labels=self.labels,
original_ttl=self.original_ttl,
expiration=self.expiration,
inception=self.inception,
key_tag=self.key_tag,
).pack()
write_name(out, encode_labels(self.signer), None)
out += self.signature
@classmethod
def unpack(cls, buf: bytes | bytearray, off: int, length: int, ctx: dict[str, Any]) -> "RRSIG":
end = off + length
fixed, pos = RRSIGFixed.unpack_from(buf, off)
signer, pos = decode_name(buf, pos)
signature = bytes(buf[pos:end])
return cls(
fixed.type_covered,
fixed.algorithm,
fixed.labels,
fixed.original_ttl,
fixed.expiration,
fixed.inception,
fixed.key_tag,
signer,
signature,
)
def encode_type_bitmap(rrtypes: list[int]) -> bytes:
"""RFC 4034 section 4.1.2: types are grouped into 256-wide windows, each
written as (window number, bitmap length in bytes, bitmap) with trailing
all-zero bytes dropped from that window's bitmap."""
windows: dict[int, bytearray] = {}
for rrtype in rrtypes:
window, bit = divmod(int(rrtype), 256)
windows.setdefault(window, bytearray(32))[bit // 8] |= 0x80 >> (bit % 8)
out = bytearray()
for window in sorted(windows):
bitmap = windows[window]
length = max(i for i, b in enumerate(bitmap) if b) + 1
out.append(window)
out.append(length)
out += bitmap[:length]
return bytes(out)
def decode_type_bitmap(buf: bytes | bytearray, off: int, end: int) -> list[int]:
rrtypes = []
pos = off
while pos < end:
window = buf[pos]
length = buf[pos + 1]
pos += 2
for i, byte in enumerate(buf[pos : pos + length]):
for bit in range(8):
if byte & (0x80 >> bit):
rrtypes.append(window * 256 + i * 8 + bit)
pos += length
return rrtypes
@dataclass
class NSEC:
"""RFC 4034 section 4; `next_domain` must not be compressed."""
RTYPE = RRType.NSEC
next_domain: str
rrtypes: list[int] = field(default_factory=list)
def pack_into(self, out: bytearray, ctx: dict[str, Any] | None) -> None:
write_name(out, encode_labels(self.next_domain), None)
out += encode_type_bitmap(self.rrtypes)
@classmethod
def unpack(cls, buf: bytes | bytearray, off: int, length: int, ctx: dict[str, Any]) -> "NSEC":
end = off + length
next_domain, pos = decode_name(buf, off)
return cls(next_domain, decode_type_bitmap(buf, pos, end))
@dataclass
class RP:
"""RFC 1183 section 2.2; both names are name-compressible, same as
CNAME/NS/PTR."""
RTYPE = RRType.RP
mbox: str
txt: str
def pack_into(self, out: bytearray, ctx: dict[str, Any] | None) -> None:
write_name(out, encode_labels(self.mbox), ctx)
write_name(out, encode_labels(self.txt), ctx)
@classmethod
def unpack(cls, buf: bytes | bytearray, off: int, length: int, ctx: dict[str, Any]) -> "RP":
mbox, pos = decode_name(buf, off)
txt, _ = decode_name(buf, pos)
return cls(mbox, txt)
class HTTPSPrefix(Struct, byteorder="big"):
"""HTTPS's one fixed-width field, ahead of `target`."""
priority: U16
class SvcParamHeader(Struct, byteorder="big"):
"""One SvcParam's (SvcParamKey, SvcParamValue length) pair, ahead of
that many bytes of raw value."""
key: U16
length: U16
@dataclass
class HTTPS:
"""RFC 9460; `target` (the SVCB TargetName) must not be compressed.
`params` is the raw (SvcParamKey, SvcParamValue) list -- key-specific
formatting (alpn/ipv4hint/echconfig/...) is a presentation-format
concern this module does not implement, same scope as the rest of this
hand-written RDATA layer."""
RTYPE = RRType.HTTPS
priority: int
target: str
params: list[tuple[int, bytes]] = field(default_factory=list)
def pack_into(self, out: bytearray, ctx: dict[str, Any] | None) -> None:
out += HTTPSPrefix(priority=self.priority).pack()
write_name(out, encode_labels(self.target), None)
for key, value in self.params:
out += SvcParamHeader(key=key, length=len(value)).pack()
out += value
@classmethod
def unpack(cls, buf: bytes | bytearray, off: int, length: int, ctx: dict[str, Any]) -> "HTTPS":
end = off + length
prefix, pos = HTTPSPrefix.unpack_from(buf, off)
target, pos = decode_name(buf, pos)
params = []
while pos < end:
header, pos = SvcParamHeader.unpack_from(buf, pos)
params.append((header.key, bytes(buf[pos : pos + header.length])))
pos += header.length
return cls(prefix.priority, target, params)
[docs]
@dataclass
class UnknownRData:
"""Raw RDATA of a record type the registry does not know."""
rtype: RRType
data: bytes
def pack_into(self, out: bytearray, ctx: dict[str, Any] | None) -> None:
out += self.data
class RawRDataClass(Protocol):
"""The non-`Struct` RDATA classes (NS/CNAME/PTR/MX/SOA/TXT): each needs
the wire's byte offset/record length/name-compression ctx to decode
(for compression-pointer resolution), unlike `Struct`'s own
self-contained `unpack(buf)` -- so they implement this 4-arg classmethod
instead of inheriting from `Struct`."""
@classmethod
def unpack(cls, buf: bytes | bytearray, off: int, length: int, ctx: dict[str, Any]) -> Any: ...
class RDataKind(enum.Enum):
"""The RDATA type registry: an enum member per implemented record type,
keyed by its wire `RRType`, so `RDataKind(rtype)` is itself the lookup
(the usual `ValueError`-on-unknown-value enum contract, no hand-rolled
dict needed). A wire rtype with no member here falls back to
`UnknownRData`."""
rdata_cls: "type[Struct] | type[RawRDataClass]"
def __new__(cls, rtype: int, rdata_cls: "type[Struct] | type[RawRDataClass]") -> "RDataKind":
obj = object.__new__(cls)
obj._value_ = rtype
obj.rdata_cls = rdata_cls
return obj
A = (RRType.A, A)
AAAA = (RRType.AAAA, AAAA)
NS = (RRType.NS, NS)
CNAME = (RRType.CNAME, CNAME)
DNAME = (RRType.DNAME, DNAME)
PTR = (RRType.PTR, PTR)
MX = (RRType.MX, MX)
SOA = (RRType.SOA, SOA)
SRV = (RRType.SRV, SRV)
NAPTR = (RRType.NAPTR, NAPTR)
DS = (RRType.DS, DS)
DNSKEY = (RRType.DNSKEY, DNSKEY)
RRSIG = (RRType.RRSIG, RRSIG)
NSEC = (RRType.NSEC, NSEC)
CAA = (RRType.CAA, CAA)
HTTPS = (RRType.HTTPS, HTTPS)
LOC = (RRType.LOC, LOC)
RP = (RRType.RP, RP)
TLSA = (RRType.TLSA, TLSA)
SSHFP = (RRType.SSHFP, SSHFP)
TXT = (RRType.TXT, TXT)
def pack_rdata(data: Any, out: bytearray, ctx: dict[str, Any] | None) -> None:
if isinstance(data, Struct):
out += data.pack()
else:
data.pack_into(out, ctx)
def unpack_rdata(rtype: int, buf: bytes | bytearray, off: int, length: int, ctx: dict[str, Any]) -> Any:
try:
rdata_cls = RDataKind(rtype).rdata_cls
except ValueError:
return UnknownRData(RRType(rtype), bytes(buf[off : off + length]))
if issubclass(rdata_cls, Struct):
return rdata_cls.unpack(bytes(buf[off : off + length]))
return rdata_cls.unpack(buf, off, length, ctx)
# ---------- Question / ResourceRecord / DNS: hand-written orchestration ----------
class QuestionTail(Struct, byteorder="big"):
"""A question's fixed-width part, after its (possibly compressed) name."""
qtype: U16
qclass: U16
[docs]
@dataclass
class Question:
name: str
qtype: RRType = RRType.A
qclass: DNSClass = DNSClass.IN
def pack_into(self, out: bytearray, ctx: dict[str, Any] | None) -> None:
write_name(out, encode_labels(self.name), ctx)
out += QuestionTail(qtype=self.qtype, qclass=self.qclass).pack()
@classmethod
def unpack(cls, buf: bytes | bytearray, off: int, ctx: dict[str, Any]) -> tuple["Question", int]:
name, off = decode_name(buf, off)
tail, pos = QuestionTail.unpack_from(buf, off)
return cls(name, RRType(tail.qtype), DNSClass(tail.qclass)), pos
class ResourceRecordHeader(Struct, byteorder="big"):
"""TYPE/CLASS/TTL, ahead of RDLENGTH and the RDATA body. Packed with
`rdlength=0` as a placeholder (see `ResourceRecord.pack_into`'s
docstring for why RDLENGTH can't be known yet at this point);
`unpack_from()` instead reads it for real, since on read the whole
record -- RDATA included -- is already in `buf`."""
rtype: U16
rclass: U16
ttl: U32
rdlength: U16 = described(default=0, help="always 0 on pack; patched in afterwards")
[docs]
@dataclass
class ResourceRecord:
"""`rtype` is never stored here: a known record type carries its own
class-level `RTYPE`, and `UnknownRData` already carries the wire tag it
was decoded with -- storing a third, possibly-inconsistent copy on the
record itself would just invite them drifting apart.
`rclass` is `DNSClass | int`, not just `DNSClass`: an EDNS0 OPT
pseudo-record (see `edns0()`) repurposes this field to carry the
requestor's/responder's UDP payload size instead of an actual class."""
name: str
rclass: DNSClass | int = DNSClass.IN
ttl: int = 0
data: object = None
def pack_into(self, out: bytearray, ctx: dict[str, Any] | None) -> None:
write_name(out, encode_labels(self.name), ctx)
rtype = getattr(self.data, "RTYPE", None)
if rtype is None:
# No class-level RTYPE means this wasn't a registered rdata_cls,
# so unpack_rdata() must have produced an UnknownRData instead --
# the only other shape `data` can take (see the class docstring).
assert isinstance(self.data, UnknownRData)
rtype = self.data.rtype
# RDLENGTH isn't known until the RDATA body itself is packed (its
# own length can depend on name compression state further down the
# message), so the whole header is written once as a 0-rdlength
# placeholder and then overwritten in place, rdlength and all.
rdlen_pos = len(out)
out += ResourceRecordHeader(rtype=rtype, rclass=self.rclass, ttl=self.ttl).pack()
body_start = len(out)
pack_rdata(self.data, out, ctx)
rdlength = len(out) - body_start
out[rdlen_pos:body_start] = ResourceRecordHeader(
rtype=rtype, rclass=self.rclass, ttl=self.ttl, rdlength=rdlength
).pack()
@classmethod
def unpack(cls, buf: bytes | bytearray, off: int, ctx: dict[str, Any]) -> tuple["ResourceRecord", int]:
name, off = decode_name(buf, off)
header, off = ResourceRecordHeader.unpack_from(buf, off)
data = unpack_rdata(header.rtype, buf, off, header.rdlength, ctx)
off += header.rdlength
return cls(name=name, rclass=DNSClass(header.rclass), ttl=header.ttl, data=data), off
class MessageHeader(Struct, byteorder="big"):
"""The whole 12-byte DNS header (RFC 1035 section 4.1.1): ID, the
FLAGS word (nested -- see `DNSFlags`) and the four section counts."""
id: U16
flags: DNSFlags
qdcount: U16
ancount: U16
nscount: U16
arcount: U16
[docs]
@dataclass
class DNS:
"""A whole DNS message; the section counts are computed on pack()."""
id: int = 0
flags: DNSFlags = field(default_factory=DNSFlags)
questions: list[Question] = field(default_factory=list)
answers: list[ResourceRecord] = field(default_factory=list)
authorities: list[ResourceRecord] = field(default_factory=list)
additionals: list[ResourceRecord] = field(default_factory=list)
def pack(self, compress: bool = True) -> bytes:
out = bytearray(
MessageHeader(
id=self.id,
flags=self.flags,
qdcount=len(self.questions),
ancount=len(self.answers),
nscount=len(self.authorities),
arcount=len(self.additionals),
).pack()
)
ctx: dict[str, Any] | None = {} if compress else None
for q in self.questions:
q.pack_into(out, ctx)
for r in self.answers:
r.pack_into(out, ctx)
for r in self.authorities:
r.pack_into(out, ctx)
for r in self.additionals:
r.pack_into(out, ctx)
return bytes(out)
@classmethod
def unpack(cls, buf: bytes | bytearray) -> "DNS":
header, off = MessageHeader.unpack_from(buf, 0)
ctx: dict[str, Any] = {}
def read_n(n: int, reader: Any) -> list[Any]:
nonlocal off
items = []
for _ in range(n):
item, off = reader(buf, off, ctx)
items.append(item)
return items
questions = read_n(header.qdcount, Question.unpack)
answers = read_n(header.ancount, ResourceRecord.unpack)
authorities = read_n(header.nscount, ResourceRecord.unpack)
additionals = read_n(header.arcount, ResourceRecord.unpack)
return cls(header.id, header.flags, questions, answers, authorities, additionals)
def reply(request: DNS, *, aa: bool = True) -> DNS:
"""A ready-to-answer `DNS` for `request`: same id and questions, `qr`
set, `rd` echoed back, `ra` left False (this module has no resolver of
its own to promise recursion from) and `rcode` NOERROR -- the caller
fills in answers/authorities/additionals and overwrites
`.flags.rcode`/`.flags.tc` on the way out, same as any other mutable
`Struct`/dataclass field."""
return DNS(
id=request.id,
flags=DNSFlags(
qr=True, opcode=request.flags.opcode, aa=aa, rd=request.flags.rd, ra=False, rcode=RCODE.NOERROR
),
questions=list(request.questions),
)
def edns0(
udp_payload_size: int,
*,
extended_rcode: int = 0,
version: int = 0,
do: bool = False,
options: bytes = b"",
) -> ResourceRecord:
"""An EDNS0 OPT pseudo-record (RFC 6891) ready to append to a message's
`additionals`. The requestor's/responder's UDP payload size is smuggled
in the ordinary `rclass` field (that's the whole point of `DNSClass`
being an `OpenIntEnum` rather than a closed set); `ttl` carries the
extended RCODE, version and the DO bit instead of an actual TTL."""
ttl = (extended_rcode << 24) | (version << 16) | (0x8000 if do else 0)
return ResourceRecord(name="", rclass=udp_payload_size, ttl=ttl, data=UnknownRData(RRType.OPT, bytes(options)))
def edns_udp_size(msg: DNS) -> int | None:
"""The requestor's/responder's advertised UDP payload size from `msg`'s
first OPT pseudo-record in `additionals`, or None if it has none."""
for rr in msg.additionals:
rtype = getattr(rr.data, "RTYPE", None)
if rtype is None and isinstance(rr.data, UnknownRData):
rtype = rr.data.rtype
if rtype == RRType.OPT:
return int(rr.rclass)
return None