2025-12-25 upload

venv/Lib/site-packages/mitmproxy/proxy/layers/dns.py

@@ -0,0 +1,190 @@
+import struct
+import time
+from dataclasses import dataclass
+from typing import List
+from typing import Literal
+
+from mitmproxy import dns
+from mitmproxy import flow as mflow
+from mitmproxy.net.dns import response_codes
+from mitmproxy.proxy import commands
+from mitmproxy.proxy import events
+from mitmproxy.proxy import layer
+from mitmproxy.proxy.context import Context
+from mitmproxy.proxy.utils import expect
+
+_LENGTH_LABEL = struct.Struct("!H")
+
+
+@dataclass
+class DnsRequestHook(commands.StartHook):
+    """
+    A DNS query has been received.
+    """
+
+    flow: dns.DNSFlow
+
+
+@dataclass
+class DnsResponseHook(commands.StartHook):
+    """
+    A DNS response has been received or set.
+    """
+
+    flow: dns.DNSFlow
+
+
+@dataclass
+class DnsErrorHook(commands.StartHook):
+    """
+    A DNS error has occurred.
+    """
+
+    flow: dns.DNSFlow
+
+
+def pack_message(
+    message: dns.DNSMessage, transport_protocol: Literal["tcp", "udp"]
+) -> bytes:
+    packed = message.packed
+    if transport_protocol == "tcp":
+        return struct.pack("!H", len(packed)) + packed
+    else:
+        return packed
+
+
+class DNSLayer(layer.Layer):
+    """
+    Layer that handles resolving DNS queries.
+    """
+
+    flows: dict[int, dns.DNSFlow]
+    req_buf: bytearray
+    resp_buf: bytearray
+
+    def __init__(self, context: Context):
+        super().__init__(context)
+        self.flows = {}
+        self.req_buf = bytearray()
+        self.resp_buf = bytearray()
+
+    def handle_request(
+        self, flow: dns.DNSFlow, msg: dns.DNSMessage
+    ) -> layer.CommandGenerator[None]:
+        flow.request = msg  # if already set, continue and query upstream again
+        yield DnsRequestHook(flow)
+        if flow.response:
+            yield from self.handle_response(flow, flow.response)
+        elif flow.error:
+            yield from self.handle_error(flow, flow.error.msg)
+        elif not self.context.server.address:
+            yield from self.handle_error(
+                flow, "No hook has set a response and there is no upstream server."
+            )
+        else:
+            if not self.context.server.connected:
+                err = yield commands.OpenConnection(self.context.server)
+                if err:
+                    yield from self.handle_error(flow, str(err))
+                    # cannot recover from this
+                    return
+            packed = pack_message(flow.request, flow.server_conn.transport_protocol)
+            yield commands.SendData(self.context.server, packed)
+
+    def handle_response(
+        self, flow: dns.DNSFlow, msg: dns.DNSMessage
+    ) -> layer.CommandGenerator[None]:
+        flow.response = msg
+        yield DnsResponseHook(flow)
+        if flow.response:
+            packed = pack_message(flow.response, flow.client_conn.transport_protocol)
+            yield commands.SendData(self.context.client, packed)
+
+    def handle_error(self, flow: dns.DNSFlow, err: str) -> layer.CommandGenerator[None]:
+        flow.error = mflow.Error(err)
+        yield DnsErrorHook(flow)
+        servfail = flow.request.fail(response_codes.SERVFAIL)
+        yield commands.SendData(
+            self.context.client,
+            pack_message(servfail, flow.client_conn.transport_protocol),
+        )
+
+    def unpack_message(self, data: bytes, from_client: bool) -> List[dns.DNSMessage]:
+        msgs: List[dns.DNSMessage] = []
+
+        buf = self.req_buf if from_client else self.resp_buf
+
+        if self.context.client.transport_protocol == "udp":
+            msgs.append(dns.DNSMessage.unpack(data, timestamp=time.time()))
+        elif self.context.client.transport_protocol == "tcp":
+            buf.extend(data)
+            size = len(buf)
+            offset = 0
+
+            while True:
+                if size - offset < _LENGTH_LABEL.size:
+                    break
+                (expected_size,) = _LENGTH_LABEL.unpack_from(buf, offset)
+                offset += _LENGTH_LABEL.size
+                if expected_size == 0:
+                    raise struct.error("Message length field cannot be zero")
+
+                if size - offset < expected_size:
+                    offset -= _LENGTH_LABEL.size
+                    break
+
+                data = bytes(buf[offset : expected_size + offset])
+                offset += expected_size
+                msgs.append(dns.DNSMessage.unpack(data, timestamp=time.time()))
+
+            del buf[:offset]
+        return msgs
+
+    @expect(events.Start)
+    def state_start(self, _) -> layer.CommandGenerator[None]:
+        self._handle_event = self.state_query
+        yield from ()
+
+    @expect(events.DataReceived, events.ConnectionClosed)
+    def state_query(self, event: events.Event) -> layer.CommandGenerator[None]:
+        assert isinstance(event, events.ConnectionEvent)
+        from_client = event.connection is self.context.client
+
+        if isinstance(event, events.DataReceived):
+            msgs: List[dns.DNSMessage] = []
+            try:
+                msgs = self.unpack_message(event.data, from_client)
+            except struct.error as e:
+                yield commands.Log(f"{event.connection} sent an invalid message: {e}")
+                yield commands.CloseConnection(event.connection)
+                self._handle_event = self.state_done
+            else:
+                for msg in msgs:
+                    try:
+                        flow = self.flows[msg.id]
+                    except KeyError:
+                        flow = dns.DNSFlow(
+                            self.context.client, self.context.server, live=True
+                        )
+                        self.flows[msg.id] = flow
+                    if from_client:
+                        yield from self.handle_request(flow, msg)
+                    else:
+                        yield from self.handle_response(flow, msg)
+
+        elif isinstance(event, events.ConnectionClosed):
+            other_conn = self.context.server if from_client else self.context.client
+            if other_conn.connected:
+                yield commands.CloseConnection(other_conn)
+            self._handle_event = self.state_done
+            for flow in self.flows.values():
+                flow.live = False
+
+        else:
+            raise AssertionError(f"Unexpected event: {event}")
+
+    @expect(events.DataReceived, events.ConnectionClosed)
+    def state_done(self, _) -> layer.CommandGenerator[None]:
+        yield from ()
+
+    _handle_event = state_start