codekingpro/portable-devtools

1from __future__ import annotations2 3import asyncio4import copy5import time6import uuid7from dataclasses import dataclass8from dataclasses import field9from typing import Any10from typing import ClassVar11 12from mitmproxy import connection13from mitmproxy import exceptions14from mitmproxy import version15from mitmproxy.coretypes import serializable16 17 18@dataclass19class Error(serializable.SerializableDataclass):20    """21    An Error.22 23    This is distinct from an protocol error response (say, a HTTP code 500),24    which is represented by a normal `mitmproxy.http.Response` object. This class is25    responsible for indicating errors that fall outside of normal protocol26    communications, like interrupted connections, timeouts, or protocol errors.27    """28 29    msg: str30    """Message describing the error."""31 32    timestamp: float = field(default_factory=time.time)33    """Unix timestamp of when this error happened."""34 35    KILLED_MESSAGE: ClassVar[str] = "Connection killed."36 37    def __str__(self):38        return self.msg39 40    def __repr__(self):41        return self.msg42 43 44class Flow(serializable.Serializable):45    """46    Base class for network flows. A flow is a collection of objects,47    for example HTTP request/response pairs or a list of TCP messages.48 49    See also:50     - mitmproxy.http.HTTPFlow51     - mitmproxy.tcp.TCPFlow52     - mitmproxy.udp.UDPFlow53    """54 55    client_conn: connection.Client56    """The client that connected to mitmproxy."""57 58    server_conn: connection.Server59    """60    The server mitmproxy connected to.61 62    Some flows may never cause mitmproxy to initiate a server connection,63    for example because their response is replayed by mitmproxy itself.64    To simplify implementation, those flows will still have a `server_conn` attribute65    with a `timestamp_start` set to `None`.66    """67 68    error: Error | None = None69    """A connection or protocol error affecting this flow."""70 71    intercepted: bool72    """73    If `True`, the flow is currently paused by mitmproxy.74    We're waiting for a user action to forward the flow to its destination.75    """76 77    marked: str = ""78    """79    If this attribute is a non-empty string the flow has been marked by the user.80 81    A string value will be used as the marker annotation. May either be a single character or a Unicode emoji name.82 83    For example `:grapes:` becomes `🍇` in views that support emoji rendering.84    Consult the [Github API Emoji List](https://api.github.com/emojis) for a list of emoji that may be used.85    Not all emoji, especially [emoji modifiers](https://en.wikipedia.org/wiki/Miscellaneous_Symbols_and_Pictographs#Emoji_modifiers)86    will render consistently.87 88    The default marker for the view will be used if the Unicode emoji name can not be interpreted.89    """90 91    is_replay: str | None92    """93    This attribute indicates if this flow has been replayed in either direction.94 95     - a value of `request` indicates that the request has been artifically replayed by mitmproxy to the server.96     - a value of `response` indicates that the response to the client's request has been set by server replay.97    """98 99    live: bool100    """101    If `True`, the flow belongs to a currently active connection.102    If `False`, the flow may have been already completed or loaded from disk.103    """104 105    timestamp_created: float106    """107    The Unix timestamp of when this flow was created.108 109    In contrast to `timestamp_start`, this value will not change when a flow is replayed.110    """111 112    def __init__(113        self,114        client_conn: connection.Client,115        server_conn: connection.Server,116        live: bool = False,117    ) -> None:118        self.id = str(uuid.uuid4())119        self.client_conn = client_conn120        self.server_conn = server_conn121        self.live = live122        self.timestamp_created = time.time()123 124        self.intercepted: bool = False125        self._resume_event: asyncio.Event | None = None126        self._backup: Flow | None = None127        self.marked: str = ""128        self.is_replay: str | None = None129        self.metadata: dict[str, Any] = dict()130        self.comment: str = ""131 132    __types: dict[str, type[Flow]] = {}133 134    type: ClassVar[135        str136    ]  # automatically derived from the class name in __init_subclass__137    """The flow type, for example `http`, `tcp`, or `dns`."""138 139    def __init_subclass__(cls, **kwargs):140        cls.type = cls.__name__.removesuffix("Flow").lower()141        Flow.__types[cls.type] = cls142 143    def get_state(self) -> serializable.State:144        state = {145            "version": version.FLOW_FORMAT_VERSION,146            "type": self.type,147            "id": self.id,148            "error": self.error.get_state() if self.error else None,149            "client_conn": self.client_conn.get_state(),150            "server_conn": self.server_conn.get_state(),151            "intercepted": self.intercepted,152            "is_replay": self.is_replay,153            "marked": self.marked,154            "metadata": copy.deepcopy(self.metadata),155            "comment": self.comment,156            "timestamp_created": self.timestamp_created,157        }158        state["backup"] = copy.deepcopy(self._backup) if self._backup != state else None159        return state160 161    def set_state(self, state: serializable.State) -> None:162        assert state.pop("version") == version.FLOW_FORMAT_VERSION163        assert state.pop("type") == self.type164        self.id = state.pop("id")165        if state["error"]:166            if self.error:167                self.error.set_state(state.pop("error"))168            else:169                self.error = Error.from_state(state.pop("error"))170        else:171            self.error = state.pop("error")172        self.client_conn.set_state(state.pop("client_conn"))173        self.server_conn.set_state(state.pop("server_conn"))174        self.intercepted = state.pop("intercepted")175        self.is_replay = state.pop("is_replay")176        self.marked = state.pop("marked")177        self.metadata = state.pop("metadata")178        self.comment = state.pop("comment")179        self.timestamp_created = state.pop("timestamp_created")180        self._backup = state.pop("backup", None)181        assert state == {}182 183    @classmethod184    def from_state(cls, state: serializable.State) -> Flow:185        try:186            flow_cls = Flow.__types[state["type"]]187        except KeyError:188            raise ValueError(f"Unknown flow type: {state['type']}")189        client = connection.Client(peername=("", 0), sockname=("", 0))190        server = connection.Server(address=None)191        f = flow_cls(client, server)192        f.set_state(state)193        return f194 195    def copy(self):196        """Make a copy of this flow."""197        f = super().copy()198        f.live = False199        return f200 201    def modified(self):202        """203        `True` if this file has been modified by a user, `False` otherwise.204        """205        if self._backup:206            return self._backup != self.get_state()207        else:208            return False209 210    def backup(self, force=False):211        """212        Save a backup of this flow, which can be restored by calling `Flow.revert()`.213        """214        if not self._backup:215            self._backup = self.get_state()216 217    def revert(self):218        """219        Revert to the last backed up state.220        """221        if self._backup:222            self.set_state(self._backup)223            self._backup = None224 225    @property226    def killable(self):227        """*Read-only:* `True` if this flow can be killed, `False` otherwise."""228        return self.live and not (self.error and self.error.msg == Error.KILLED_MESSAGE)229 230    def kill(self):231        """232        Kill this flow. The current request/response will not be forwarded to its destination.233        """234        if not self.killable:235            raise exceptions.ControlException("Flow is not killable.")236        # TODO: The way we currently signal killing is not ideal. One major problem is that we cannot kill237        #  flows in transit (https://github.com/mitmproxy/mitmproxy/issues/4711), even though they are advertised238        #  as killable. An alternative approach would be to introduce a `KillInjected` event similar to239        #  `MessageInjected`, which should fix this issue.240        self.error = Error(Error.KILLED_MESSAGE)241        self.intercepted = False242        self.live = False243 244    def intercept(self):245        """246        Intercept this Flow. Processing will stop until resume is247        called.248        """249        if self.intercepted:250            return251        self.intercepted = True252        if self._resume_event is not None:253            self._resume_event.clear()254 255    async def wait_for_resume(self):256        """257        Wait until this Flow is resumed.258        """259        if not self.intercepted:260            return261        if self._resume_event is None:262            self._resume_event = asyncio.Event()263        await self._resume_event.wait()264 265    def resume(self):266        """267        Continue with the flow – called after an intercept().268        """269        if not self.intercepted:270            return271        self.intercepted = False272        if self._resume_event is not None:273            self._resume_event.set()274 275    @property276    def timestamp_start(self) -> float:277        """278        *Read-only:* Start time of the flow.279        Depending on the flow type, this property is an alias for280        `mitmproxy.connection.Client.timestamp_start` or `mitmproxy.http.Request.timestamp_start`.281        """282        return self.client_conn.timestamp_start283 284 285__all__ = [286    "Flow",287    "Error",288]289