codekingpro/portable-devtools

1"""2When IO actions occur at the proxy server, they are passed down to layers as events.3Events represent the only way for layers to receive new data from sockets.4The counterpart to events are commands.5"""6 7import typing8import warnings9from dataclasses import dataclass10from dataclasses import is_dataclass11from typing import Any12from typing import Generic13from typing import TypeVar14 15from mitmproxy import flow16from mitmproxy.connection import Connection17from mitmproxy.proxy import commands18 19 20class Event:21    """22    Base class for all events.23    """24 25    def __repr__(self):26        return f"{type(self).__name__}({self.__dict__!r})"27 28 29class Start(Event):30    """31    Every layer initially receives a start event.32    This is useful to emit events on startup.33    """34 35 36@dataclass37class ConnectionEvent(Event):38    """39    All events involving connection IO.40    """41 42    connection: Connection43 44 45@dataclass46class DataReceived(ConnectionEvent):47    """48    Remote has sent some data.49    """50 51    data: bytes52 53    def __repr__(self):54        target = type(self.connection).__name__.lower()55        return f"DataReceived({target}, {self.data!r})"56 57 58class ConnectionClosed(ConnectionEvent):59    """60    Remote has closed a connection.61    """62 63 64class CommandCompleted(Event):65    """66    Emitted when a command has been finished, e.g.67    when the master has replied or when we have established a server connection.68    """69 70    command: commands.Command71    reply: Any72 73    def __new__(cls, *args, **kwargs):74        if cls is CommandCompleted:75            raise TypeError("CommandCompleted may not be instantiated directly.")76        assert is_dataclass(cls)77        return super().__new__(cls)78 79    def __init_subclass__(cls, **kwargs):80        command_cls = typing.get_type_hints(cls).get("command", None)81        valid_command_subclass = (82            isinstance(command_cls, type)83            and issubclass(command_cls, commands.Command)84            and command_cls is not commands.Command85        )86        if not valid_command_subclass:87            warnings.warn(88                f"{cls} needs a properly annotated command attribute.",89                RuntimeWarning,90            )91        if command_cls in command_reply_subclasses:92            other = command_reply_subclasses[command_cls]93            warnings.warn(94                f"Two conflicting subclasses for {command_cls}: {cls} and {other}",95                RuntimeWarning,96            )97        command_reply_subclasses[command_cls] = cls  # type: ignore98 99    def __repr__(self):100        return f"Reply({self.command!r}, {self.reply!r})"101 102 103command_reply_subclasses: dict[commands.Command, type[CommandCompleted]] = {}104 105 106@dataclass(repr=False)107class OpenConnectionCompleted(CommandCompleted):108    command: commands.OpenConnection109    reply: str | None110    """error message"""111 112 113@dataclass(repr=False)114class HookCompleted(CommandCompleted):115    command: commands.StartHook116    reply: None = None117 118 119T = TypeVar("T")120 121 122@dataclass123class MessageInjected(Event, Generic[T]):124    """125    The user has injected a custom WebSocket/TCP/... message.126    """127 128    flow: flow.Flow129    message: T130 131 132@dataclass133class Wakeup(CommandCompleted):134    """135    Event sent to layers that requested a wakeup using RequestWakeup.136    """137 138    command: commands.RequestWakeup139