2025-12-25 upload

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

@@ -0,0 +1,132 @@
+from dataclasses import dataclass
+
+from mitmproxy import flow
+from mitmproxy import udp
+from mitmproxy.connection import Connection
+from mitmproxy.proxy import commands
+from mitmproxy.proxy import events
+from mitmproxy.proxy import layer
+from mitmproxy.proxy.commands import StartHook
+from mitmproxy.proxy.context import Context
+from mitmproxy.proxy.events import MessageInjected
+from mitmproxy.proxy.utils import expect
+
+
+@dataclass
+class UdpStartHook(StartHook):
+    """
+    A UDP connection has started.
+    """
+
+    flow: udp.UDPFlow
+
+
+@dataclass
+class UdpMessageHook(StartHook):
+    """
+    A UDP connection has received a message. The most recent message
+    will be flow.messages[-1]. The message is user-modifiable.
+    """
+
+    flow: udp.UDPFlow
+
+
+@dataclass
+class UdpEndHook(StartHook):
+    """
+    A UDP connection has ended.
+    """
+
+    flow: udp.UDPFlow
+
+
+@dataclass
+class UdpErrorHook(StartHook):
+    """
+    A UDP error has occurred.
+
+    Every UDP flow will receive either a udp_error or a udp_end event, but not both.
+    """
+
+    flow: udp.UDPFlow
+
+
+class UdpMessageInjected(MessageInjected[udp.UDPMessage]):
+    """
+    The user has injected a custom UDP message.
+    """
+
+
+class UDPLayer(layer.Layer):
+    """
+    Simple UDP layer that just relays messages right now.
+    """
+
+    flow: udp.UDPFlow | None
+
+    def __init__(self, context: Context, ignore: bool = False):
+        super().__init__(context)
+        if ignore:
+            self.flow = None
+        else:
+            self.flow = udp.UDPFlow(self.context.client, self.context.server, True)
+
+    @expect(events.Start)
+    def start(self, _) -> layer.CommandGenerator[None]:
+        if self.flow:
+            yield UdpStartHook(self.flow)
+
+        if self.context.server.timestamp_start is None:
+            err = yield commands.OpenConnection(self.context.server)
+            if err:
+                if self.flow:
+                    self.flow.error = flow.Error(str(err))
+                    yield UdpErrorHook(self.flow)
+                yield commands.CloseConnection(self.context.client)
+                self._handle_event = self.done
+                return
+        self._handle_event = self.relay_messages
+
+    _handle_event = start
+
+    @expect(events.DataReceived, events.ConnectionClosed, UdpMessageInjected)
+    def relay_messages(self, event: events.Event) -> layer.CommandGenerator[None]:
+        if isinstance(event, UdpMessageInjected):
+            # we just spoof that we received data here and then process that regularly.
+            event = events.DataReceived(
+                self.context.client
+                if event.message.from_client
+                else self.context.server,
+                event.message.content,
+            )
+
+        assert isinstance(event, events.ConnectionEvent)
+
+        from_client = event.connection == self.context.client
+        send_to: Connection
+        if from_client:
+            send_to = self.context.server
+        else:
+            send_to = self.context.client
+
+        if isinstance(event, events.DataReceived):
+            if self.flow:
+                udp_message = udp.UDPMessage(from_client, event.data)
+                self.flow.messages.append(udp_message)
+                yield UdpMessageHook(self.flow)
+                yield commands.SendData(send_to, udp_message.content)
+            else:
+                yield commands.SendData(send_to, event.data)
+
+        elif isinstance(event, events.ConnectionClosed):
+            self._handle_event = self.done
+            yield commands.CloseConnection(send_to)
+            if self.flow:
+                yield UdpEndHook(self.flow)
+                self.flow.live = False
+        else:
+            raise AssertionError(f"Unexpected event: {event}")
+
+    @expect(events.DataReceived, events.ConnectionClosed, UdpMessageInjected)
+    def done(self, _) -> layer.CommandGenerator[None]:
+        yield from ()