133 lines
3.8 KiB
Python
133 lines
3.8 KiB
Python
|
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 ()
|