codekingpro/portable-devtools

1from __future__ import annotations2 3from mitmproxy import connection4from mitmproxy.proxy import commands5 6 7class QuicStreamCommand(commands.ConnectionCommand):8    """Base class for all QUIC stream commands."""9 10    stream_id: int11    """The ID of the stream the command was issued for."""12 13    def __init__(self, connection: connection.Connection, stream_id: int) -> None:14        super().__init__(connection)15        self.stream_id = stream_id16 17 18class SendQuicStreamData(QuicStreamCommand):19    """Command that sends data on a stream."""20 21    data: bytes22    """The data which should be sent."""23    end_stream: bool24    """Whether the FIN bit should be set in the STREAM frame."""25 26    def __init__(27        self,28        connection: connection.Connection,29        stream_id: int,30        data: bytes,31        end_stream: bool = False,32    ) -> None:33        super().__init__(connection, stream_id)34        self.data = data35        self.end_stream = end_stream36 37    def __repr__(self):38        target = repr(self.connection).partition("(")[0].lower()39        end_stream = "[end_stream] " if self.end_stream else ""40        return f"SendQuicStreamData({target} on {self.stream_id}, {end_stream}{self.data!r})"41 42 43class ResetQuicStream(QuicStreamCommand):44    """Abruptly terminate the sending part of a stream."""45 46    error_code: int47    """An error code indicating why the stream is being reset."""48 49    def __init__(50        self, connection: connection.Connection, stream_id: int, error_code: int51    ) -> None:52        super().__init__(connection, stream_id)53        self.error_code = error_code54 55 56class StopSendingQuicStream(QuicStreamCommand):57    """Request termination of the receiving part of a stream."""58 59    error_code: int60    """An error code indicating why the stream is being stopped."""61 62    def __init__(63        self, connection: connection.Connection, stream_id: int, error_code: int64    ) -> None:65        super().__init__(connection, stream_id)66        self.error_code = error_code67 68 69class CloseQuicConnection(commands.CloseConnection):70    """Close a QUIC connection."""71 72    error_code: int73    "The error code which was specified when closing the connection."74 75    frame_type: int | None76    "The frame type which caused the connection to be closed, or `None`."77 78    reason_phrase: str79    "The human-readable reason for which the connection was closed."80 81    # XXX: A bit much boilerplate right now. Should switch to dataclasses.82    def __init__(83        self,84        conn: connection.Connection,85        error_code: int,86        frame_type: int | None,87        reason_phrase: str,88    ) -> None:89        super().__init__(conn)90        self.error_code = error_code91        self.frame_type = frame_type92        self.reason_phrase = reason_phrase93