codekingpro/portable-devtools

1"""2Commands make it possible for layers to communicate with the "outer world",3e.g. to perform IO or to ask the master.4A command is issued by a proxy layer and is then passed upwards to the proxy server, and from there5possibly to the master and addons.6 7The counterpart to commands are events.8"""9 10import logging11import warnings12from typing import TYPE_CHECKING13from typing import Union14 15import mitmproxy.hooks16from mitmproxy.connection import Connection17from mitmproxy.connection import Server18 19if TYPE_CHECKING:20    import mitmproxy.proxy.layer21 22 23class Command:24    """25    Base class for all commands26    """27 28    blocking: Union[bool, "mitmproxy.proxy.layer.Layer"] = False29    """30    Determines if the command blocks until it has been completed.31    For practical purposes, this attribute should be thought of as a boolean value,32    layers may swap out `True` with a reference to themselves to signal to outer layers33    that they do not need to block as well.34 35    Example:36 37        reply = yield Hook("requestheaders", flow)  # blocking command38        yield Log("hello world", "info")            # non-blocking39    """40 41    def __repr__(self):42        x = self.__dict__.copy()43        x.pop("blocking", None)44        return f"{type(self).__name__}({x!r})"45 46 47class RequestWakeup(Command):48    """49    Request a `Wakeup` event after the specified amount of seconds.50    """51 52    delay: float53 54    def __init__(self, delay: float):55        self.delay = delay56 57 58class ConnectionCommand(Command):59    """60    Commands involving a specific connection61    """62 63    connection: Connection64 65    def __init__(self, connection: Connection):66        self.connection = connection67 68 69class SendData(ConnectionCommand):70    """71    Send data to a remote peer72    """73 74    data: bytes75 76    def __init__(self, connection: Connection, data: bytes):77        super().__init__(connection)78        self.data = data79 80    def __repr__(self):81        target = str(self.connection).split("(", 1)[0].lower()82        return f"SendData({target}, {self.data!r})"83 84 85class OpenConnection(ConnectionCommand):86    """87    Open a new connection88    """89 90    connection: Server91    blocking = True92 93 94class CloseConnection(ConnectionCommand):95    """96    Close a connection. If the client connection is closed,97    all other connections will ultimately be closed during cleanup.98    """99 100 101class CloseTcpConnection(CloseConnection):102    half_close: bool103    """104    If True, only close our half of the connection by sending a FIN packet.105    This is required from some protocols which close their end to signal completion and then continue reading,106    for example HTTP/1.0 without Content-Length header.107    """108 109    def __init__(self, connection: Connection, half_close: bool = False):110        super().__init__(connection)111        self.half_close = half_close112 113 114class StartHook(Command, mitmproxy.hooks.Hook):115    """116    Start an event hook in the mitmproxy core.117    This triggers a particular function (derived from the class name) in all addons.118    """119 120    name = ""121    blocking = True122 123    def __new__(cls, *args, **kwargs):124        if cls is StartHook:125            raise TypeError("StartHook may not be instantiated directly.")126        return super().__new__(cls, *args, **kwargs)127 128 129class Log(Command):130    """131    Log a message.132 133    Layers could technically call `logging.log` directly, but the use of a command allows us to134    write more expressive playbook tests. Put differently, by using commands we can assert that135    a specific log message is a direct consequence of a particular I/O event.136    This could also be implemented with some more playbook magic in the future,137    but for now we keep the current approach as the fully sans-io one.138    """139 140    message: str141    level: int142 143    def __init__(144        self,145        message: str,146        level: int = logging.INFO,147    ):148        if isinstance(level, str):  # pragma: no cover149            warnings.warn(150                "commands.Log() now expects an integer log level, not a string.",151                DeprecationWarning,152                stacklevel=2,153            )154            level = getattr(logging, level.upper())155        self.message = message156        self.level = level157 158    def __repr__(self):159        return f"Log({self.message!r}, {logging.getLevelName(self.level).lower()})"160