mitmproxy.flow

  1import asyncio
  2import time
  3import uuid
  4from typing import Any, ClassVar, Optional
  5
  6from mitmproxy import connection
  7from mitmproxy import exceptions
  8from mitmproxy import stateobject
  9from mitmproxy import version
 10
 11
 12class Error(stateobject.StateObject):
 13    """
 14    An Error.
 15
 16    This is distinct from an protocol error response (say, a HTTP code 500),
 17    which is represented by a normal `mitmproxy.http.Response` object. This class is
 18    responsible for indicating errors that fall outside of normal protocol
 19    communications, like interrupted connections, timeouts, or protocol errors.
 20    """
 21
 22    msg: str
 23    """Message describing the error."""
 24
 25    timestamp: float
 26    """Unix timestamp of when this error happened."""
 27
 28    KILLED_MESSAGE: ClassVar[str] = "Connection killed."
 29
 30    def __init__(self, msg: str, timestamp: Optional[float] = None) -> None:
 31        """Create an error. If no timestamp is passed, the current time is used."""
 32        self.msg = msg
 33        self.timestamp = timestamp or time.time()
 34
 35    _stateobject_attributes = dict(msg=str, timestamp=float)
 36
 37    def __str__(self):
 38        return self.msg
 39
 40    def __repr__(self):
 41        return self.msg
 42
 43    @classmethod
 44    def from_state(cls, state):
 45        # the default implementation assumes an empty constructor. Override
 46        # accordingly.
 47        f = cls(None)
 48        f.set_state(state)
 49        return f
 50
 51
 52class Flow(stateobject.StateObject):
 53    """
 54    Base class for network flows. A flow is a collection of objects,
 55    for example HTTP request/response pairs or a list of TCP messages.
 56
 57    See also:
 58     - mitmproxy.http.HTTPFlow
 59     - mitmproxy.tcp.TCPFlow
 60     - mitmproxy.udp.UDPFlow
 61    """
 62
 63    client_conn: connection.Client
 64    """The client that connected to mitmproxy."""
 65
 66    server_conn: connection.Server
 67    """
 68    The server mitmproxy connected to.
 69
 70    Some flows may never cause mitmproxy to initiate a server connection,
 71    for example because their response is replayed by mitmproxy itself.
 72    To simplify implementation, those flows will still have a `server_conn` attribute
 73    with a `timestamp_start` set to `None`.
 74    """
 75
 76    error: Optional[Error] = None
 77    """A connection or protocol error affecting this flow."""
 78
 79    intercepted: bool
 80    """
 81    If `True`, the flow is currently paused by mitmproxy.
 82    We're waiting for a user action to forward the flow to its destination.
 83    """
 84
 85    marked: str = ""
 86    """
 87    If this attribute is a non-empty string the flow has been marked by the user.
 88
 89    A string value will be used as the marker annotation. May either be a single character or a Unicode emoji name.
 90
 91    For example `:grapes:` becomes `🍇` in views that support emoji rendering.
 92    Consult the [Github API Emoji List](https://api.github.com/emojis) for a list of emoji that may be used.
 93    Not all emoji, especially [emoji modifiers](https://en.wikipedia.org/wiki/Miscellaneous_Symbols_and_Pictographs#Emoji_modifiers)
 94    will render consistently.
 95
 96    The default marker for the view will be used if the Unicode emoji name can not be interpreted.
 97    """
 98
 99    is_replay: Optional[str]
100    """
101    This attribute indicates if this flow has been replayed in either direction.
102
103     - a value of `request` indicates that the request has been artifically replayed by mitmproxy to the server.
104     - a value of `response` indicates that the response to the client's request has been set by server replay.
105    """
106
107    live: bool
108    """
109    If `True`, the flow belongs to a currently active connection.
110    If `False`, the flow may have been already completed or loaded from disk.
111    """
112
113    timestamp_created: float
114    """
115    The Unix timestamp of when this flow was created.
116
117    In contrast to `timestamp_start`, this value will not change when a flow is replayed.
118    """
119
120    def __init__(
121        self,
122        client_conn: connection.Client,
123        server_conn: connection.Server,
124        live: bool = False,
125    ) -> None:
126        self.id = str(uuid.uuid4())
127        self.client_conn = client_conn
128        self.server_conn = server_conn
129        self.live = live
130        self.timestamp_created = time.time()
131
132        self.intercepted: bool = False
133        self._resume_event: Optional[asyncio.Event] = None
134        self._backup: Optional[Flow] = None
135        self.marked: str = ""
136        self.is_replay: Optional[str] = None
137        self.metadata: dict[str, Any] = dict()
138        self.comment: str = ""
139
140    _stateobject_attributes = dict(
141        id=str,
142        error=Error,
143        client_conn=connection.Client,
144        server_conn=connection.Server,
145        intercepted=bool,
146        is_replay=str,
147        marked=str,
148        metadata=dict[str, Any],
149        comment=str,
150        timestamp_created=float,
151    )
152
153    __types: dict[str, type["Flow"]] = {}
154
155    type: ClassVar[str]  # automatically derived from the class name in __init_subclass__
156    """The flow type, for example `http`, `tcp`, or `dns`."""
157
158    def __init_subclass__(cls, **kwargs):
159        cls.type = cls.__name__.removesuffix("Flow").lower()
160        Flow.__types[cls.type] = cls
161
162    def get_state(self):
163        d = super().get_state()
164        d.update(version=version.FLOW_FORMAT_VERSION, type=self.type)
165        if self._backup and self._backup != d:
166            d.update(backup=self._backup)
167        return d
168
169    def set_state(self, state):
170        state = state.copy()
171        state.pop("version")
172        state.pop("type")
173        if "backup" in state:
174            self._backup = state.pop("backup")
175        super().set_state(state)
176
177    @classmethod
178    def from_state(cls, state):
179        try:
180            flow_cls = Flow.__types[state["type"]]
181        except KeyError:
182            raise ValueError(f"Unknown flow type: {state['type']}")
183        f = flow_cls(None, None)  # noqa
184        f.set_state(state)
185        return f
186
187    def copy(self):
188        """Make a copy of this flow."""
189        f = super().copy()
190        f.live = False
191        return f
192
193    def modified(self):
194        """
195        `True` if this file has been modified by a user, `False` otherwise.
196        """
197        if self._backup:
198            return self._backup != self.get_state()
199        else:
200            return False
201
202    def backup(self, force=False):
203        """
204        Save a backup of this flow, which can be restored by calling `Flow.revert()`.
205        """
206        if not self._backup:
207            self._backup = self.get_state()
208
209    def revert(self):
210        """
211        Revert to the last backed up state.
212        """
213        if self._backup:
214            self.set_state(self._backup)
215            self._backup = None
216
217    @property
218    def killable(self):
219        """*Read-only:* `True` if this flow can be killed, `False` otherwise."""
220        return self.live and not (self.error and self.error.msg == Error.KILLED_MESSAGE)
221
222    def kill(self):
223        """
224        Kill this flow. The current request/response will not be forwarded to its destination.
225        """
226        if not self.killable:
227            raise exceptions.ControlException("Flow is not killable.")
228        # TODO: The way we currently signal killing is not ideal. One major problem is that we cannot kill
229        #  flows in transit (https://github.com/mitmproxy/mitmproxy/issues/4711), even though they are advertised
230        #  as killable. An alternative approach would be to introduce a `KillInjected` event similar to
231        #  `MessageInjected`, which should fix this issue.
232        self.error = Error(Error.KILLED_MESSAGE)
233        self.intercepted = False
234        self.live = False
235
236    def intercept(self):
237        """
238        Intercept this Flow. Processing will stop until resume is
239        called.
240        """
241        if self.intercepted:
242            return
243        self.intercepted = True
244        if self._resume_event is not None:
245            self._resume_event.clear()
246
247    async def wait_for_resume(self):
248        """
249        Wait until this Flow is resumed.
250        """
251        if not self.intercepted:
252            return
253        if self._resume_event is None:
254            self._resume_event = asyncio.Event()
255        await self._resume_event.wait()
256
257    def resume(self):
258        """
259        Continue with the flow – called after an intercept().
260        """
261        if not self.intercepted:
262            return
263        self.intercepted = False
264        if self._resume_event is not None:
265            self._resume_event.set()
266
267    @property
268    def timestamp_start(self) -> float:
269        """
270        *Read-only:* Start time of the flow.
271        Depending on the flow type, this property is an alias for
272        `mitmproxy.connection.Client.timestamp_start` or `mitmproxy.http.Request.timestamp_start`.
273        """
274        return self.client_conn.timestamp_start
275
276
277__all__ = [
278    "Flow",
279    "Error",
280]