codekingpro/portable-devtools

1import enum2import struct3from dataclasses import dataclass4from typing import Self5 6from ...utils import strutils7from . import domain_names8 9"""10HTTPS records are formatted as follows (as per RFC9460):11- a 2-octet field for SvcPriority as an integer in network byte order.12- the uncompressed, fully qualified TargetName, represented as a sequence of length-prefixed labels per Section 3.1 of [RFC1035].13- the SvcParams, consuming the remainder of the record (so smaller than 65535 octets and constrained by the RDATA and DNS message sizes).14 15When the list of SvcParams is non-empty, it contains a series of SvcParamKey=SvcParamValue pairs, represented as:16- a 2-octet field containing the SvcParamKey as an integer in network byte order. (See Section 14.3.2 for the defined values.)17- a 2-octet field containing the length of the SvcParamValue as an integer between 0 and 65535 in network byte order.18- an octet string of this length whose contents are the SvcParamValue in a format determined by the SvcParamKey.19 20    https://datatracker.ietf.org/doc/rfc9460/21    https://datatracker.ietf.org/doc/rfc1035/22"""23 24 25class SVCParamKeys(enum.Enum):26    MANDATORY = 027    ALPN = 128    NO_DEFAULT_ALPN = 229    PORT = 330    IPV4HINT = 431    ECH = 532    IPV6HINT = 633 34 35type HTTPSRecordJSON = dict[str | int, str | int]36 37 38@dataclass39class HTTPSRecord:40    priority: int41    target_name: str42    params: dict[int, bytes]43 44    def __repr__(self):45        return str(self.to_json())46 47    def to_json(self) -> HTTPSRecordJSON:48        ret: HTTPSRecordJSON = {49            "target_name": self.target_name,50            "priority": self.priority,51        }52        typ: str | int53        for typ, val in self.params.items():54            try:55                typ = SVCParamKeys(typ).name.lower()56            except ValueError:57                pass58            ret[typ] = strutils.bytes_to_escaped_str(val)59        return ret60 61    @classmethod62    def from_json(cls, data: HTTPSRecordJSON) -> Self:63        target_name = data.pop("target_name")64        assert isinstance(target_name, str)65        priority = data.pop("priority")66        assert isinstance(priority, int)67        params: dict[int, bytes] = {}68        for k, v in data.items():69            if isinstance(k, str):70                k = SVCParamKeys[k.upper()].value71            assert isinstance(v, str)72            params[k] = strutils.escaped_str_to_bytes(v)73        return cls(target_name=target_name, priority=priority, params=params)74 75 76def _unpack_params(data: bytes, offset: int) -> dict[int, bytes]:77    """Unpacks the service parameters from the given offset."""78    params = {}79    while offset < len(data):80        param_type = struct.unpack("!H", data[offset : offset + 2])[0]81        offset += 282        param_length = struct.unpack("!H", data[offset : offset + 2])[0]83        offset += 284        if offset + param_length > len(data):85            raise struct.error(86                "unpack requires a buffer of %i bytes" % (offset + param_length)87            )88        param_value = data[offset : offset + param_length]89        offset += param_length90        params[param_type] = param_value91    return params92 93 94def unpack(data: bytes) -> HTTPSRecord:95    """96    Unpacks HTTPS RDATA from byte data.97 98    Raises:99        struct.error if the record is malformed.100    """101    offset = 0102 103    # Priority (2 bytes)104    priority = struct.unpack("!h", data[offset : offset + 2])[0]105    offset += 2106 107    # TargetName (variable length)108    target_name, offset = domain_names.unpack_from(data, offset)109 110    # Service Parameters (remaining bytes)111    params = _unpack_params(data, offset)112 113    return HTTPSRecord(priority=priority, target_name=target_name, params=params)114 115 116def _pack_params(params: dict[int, bytes]) -> bytes:117    """Converts the service parameters into the raw byte format"""118    buffer = bytearray()119 120    for k, v in params.items():121        buffer.extend(struct.pack("!H", k))122        buffer.extend(struct.pack("!H", len(v)))123        buffer.extend(v)124 125    return bytes(buffer)126 127 128def pack(record: HTTPSRecord) -> bytes:129    """Packs the HTTPS record into its bytes form."""130    buffer = bytearray()131    buffer.extend(struct.pack("!h", record.priority))132    buffer.extend(domain_names.pack(record.target_name))133    buffer.extend(_pack_params(record.params))134    return bytes(buffer)135