mitmproxy.udp

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