73 lines
1.8 KiB
Python
73 lines
1.8 KiB
Python
|
import time
|
||
|
|
|
||
|
|
from mitmproxy import connection
|
||
|
|
from mitmproxy import flow
|
||
|
|
from mitmproxy.coretypes import serializable
|
||
|
|
|
||
|
|
|
||
|
|
class UDPMessage(serializable.Serializable):
|
||
|
|
"""
|
||
|
|
An individual UDP datagram.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, from_client, content, timestamp=None):
|
||
|
|
self.from_client = from_client
|
||
|
|
self.content = content
|
||
|
|
self.timestamp = timestamp or time.time()
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_state(cls, state):
|
||
|
|
return cls(*state)
|
||
|
|
|
||
|
|
def get_state(self):
|
||
|
|
return self.from_client, self.content, self.timestamp
|
||
|
|
|
||
|
|
def set_state(self, state):
|
||
|
|
self.from_client, self.content, self.timestamp = state
|
||
|
|
|
||
|
|
def __repr__(self):
|
||
|
|
return "{direction} {content}".format(
|
||
|
|
direction="->" if self.from_client else "<-", content=repr(self.content)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class UDPFlow(flow.Flow):
|
||
|
|
"""
|
||
|
|
A UDPFlow is a representation of a UDP session.
|
||
|
|
"""
|
||
|
|
|
||
|
|
messages: list[UDPMessage]
|
||
|
|
"""
|
||
|
|
The messages transmitted over this connection.
|
||
|
|
|
||
|
|
The latest message can be accessed as `flow.messages[-1]` in event hooks.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
client_conn: connection.Client,
|
||
|
|
server_conn: connection.Server,
|
||
|
|
live: bool = False,
|
||
|
|
):
|
||
|
|
super().__init__(client_conn, server_conn, live)
|
||
|
|
self.messages = []
|
||
|
|
|
||
|
|
def get_state(self) -> serializable.State:
|
||
|
|
return {
|
||
|
|
**super().get_state(),
|
||
|
|
"messages": [m.get_state() for m in self.messages],
|
||
|
|
}
|
||
|
|
|
||
|
|
def set_state(self, state: serializable.State) -> None:
|
||
|
|
self.messages = [UDPMessage.from_state(m) for m in state.pop("messages")]
|
||
|
|
super().set_state(state)
|
||
|
|
|
||
|
|
def __repr__(self):
|
||
|
|
return f"<UDPFlow ({len(self.messages)} messages)>"
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
"UDPFlow",
|
||
|
|
"UDPMessage",
|
||
|
|
]
|