codekingpro/portable-devtools

1import time2 3from mitmproxy import connection4from mitmproxy import flow5from mitmproxy.coretypes import serializable6 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_client15        self.content = content16        self.timestamp = timestamp or time.time()17 18    @classmethod19    def from_state(cls, state):20        return cls(*state)21 22    def get_state(self):23        return self.from_client, self.content, self.timestamp24 25    def set_state(self, state):26        self.from_client, self.content, self.timestamp = state27 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]73