mitmproxy.udp

 1import time
 2
 3from mitmproxy import connection, flow
 4from mitmproxy.coretypes import serializable
 5
 6
 7class UDPMessage(serializable.Serializable):
 8    """
 9    An individual UDP datagram.
10    """
11
12    def __init__(self, from_client, content, timestamp=None):
13        self.from_client = from_client
14        self.content = content
15        self.timestamp = timestamp or time.time()
16
17    @classmethod
18    def from_state(cls, state):
19        return cls(*state)
20
21    def get_state(self):
22        return self.from_client, self.content, self.timestamp
23
24    def set_state(self, state):
25        self.from_client, self.content, self.timestamp = state
26
27    def __repr__(self):
28        return "{direction} {content}".format(
29            direction="->" if self.from_client else "<-", content=repr(self.content)
30        )
31
32
33class UDPFlow(flow.Flow):
34    """
35    A UDPFlow is a representation of a UDP session.
36    """
37
38    messages: list[UDPMessage]
39    """
40    The messages transmitted over this connection.
41
42    The latest message can be accessed as `flow.messages[-1]` in event hooks.
43    """
44
45    def __init__(
46        self,
47        client_conn: connection.Client,
48        server_conn: connection.Server,
49        live: bool = False,
50    ):
51        super().__init__(client_conn, server_conn, live)
52        self.messages = []
53
54    _stateobject_attributes = flow.Flow._stateobject_attributes.copy()
55    _stateobject_attributes["messages"] = list[UDPMessage]
56
57    def __repr__(self):
58        return f"<UDPFlow ({len(self.messages)} messages)>"
59
60
61__all__ = [
62    "UDPFlow",
63    "UDPMessage",
64]