codekingpro/portable-devtools

1from dataclasses import dataclass2 3from mitmproxy import flow4from mitmproxy import udp5from mitmproxy.connection import Connection6from mitmproxy.proxy import commands7from mitmproxy.proxy import events8from mitmproxy.proxy import layer9from mitmproxy.proxy.commands import StartHook10from mitmproxy.proxy.context import Context11from mitmproxy.proxy.events import MessageInjected12from mitmproxy.proxy.utils import expect13 14 15@dataclass16class UdpStartHook(StartHook):17    """18    A UDP connection has started.19    """20 21    flow: udp.UDPFlow22 23 24@dataclass25class UdpMessageHook(StartHook):26    """27    A UDP connection has received a message. The most recent message28    will be flow.messages[-1]. The message is user-modifiable.29    """30 31    flow: udp.UDPFlow32 33 34@dataclass35class UdpEndHook(StartHook):36    """37    A UDP connection has ended.38    """39 40    flow: udp.UDPFlow41 42 43@dataclass44class UdpErrorHook(StartHook):45    """46    A UDP error has occurred.47 48    Every UDP flow will receive either a udp_error or a udp_end event, but not both.49    """50 51    flow: udp.UDPFlow52 53 54class UdpMessageInjected(MessageInjected[udp.UDPMessage]):55    """56    The user has injected a custom UDP message.57    """58 59 60class UDPLayer(layer.Layer):61    """62    Simple UDP layer that just relays messages right now.63    """64 65    flow: udp.UDPFlow | None66 67    def __init__(self, context: Context, ignore: bool = False):68        super().__init__(context)69        if ignore:70            self.flow = None71        else:72            self.flow = udp.UDPFlow(self.context.client, self.context.server, True)73 74    @expect(events.Start)75    def start(self, _) -> layer.CommandGenerator[None]:76        if self.flow:77            yield UdpStartHook(self.flow)78 79        if self.context.server.timestamp_start is None:80            err = yield commands.OpenConnection(self.context.server)81            if err:82                if self.flow:83                    self.flow.error = flow.Error(str(err))84                    yield UdpErrorHook(self.flow)85                yield commands.CloseConnection(self.context.client)86                self._handle_event = self.done87                return88        self._handle_event = self.relay_messages89 90    _handle_event = start91 92    @expect(events.DataReceived, events.ConnectionClosed, UdpMessageInjected)93    def relay_messages(self, event: events.Event) -> layer.CommandGenerator[None]:94        if isinstance(event, UdpMessageInjected):95            # we just spoof that we received data here and then process that regularly.96            event = events.DataReceived(97                self.context.client98                if event.message.from_client99                else self.context.server,100                event.message.content,101            )102 103        assert isinstance(event, events.ConnectionEvent)104 105        from_client = event.connection == self.context.client106        send_to: Connection107        if from_client:108            send_to = self.context.server109        else:110            send_to = self.context.client111 112        if isinstance(event, events.DataReceived):113            if self.flow:114                udp_message = udp.UDPMessage(from_client, event.data)115                self.flow.messages.append(udp_message)116                yield UdpMessageHook(self.flow)117                yield commands.SendData(send_to, udp_message.content)118            else:119                yield commands.SendData(send_to, event.data)120 121        elif isinstance(event, events.ConnectionClosed):122            self._handle_event = self.done123            yield commands.CloseConnection(send_to)124            if self.flow:125                yield UdpEndHook(self.flow)126                self.flow.live = False127        else:128            raise AssertionError(f"Unexpected event: {event}")129 130    @expect(events.DataReceived, events.ConnectionClosed, UdpMessageInjected)131    def done(self, _) -> layer.CommandGenerator[None]:132        yield from ()133