codekingpro/portable-devtools

1import time2from enum import auto3from enum import Enum4from typing import Union5 6from mitmproxy import connection7from mitmproxy.proxy import commands8from mitmproxy.proxy import context9from mitmproxy.proxy import events10from mitmproxy.proxy import layer11from mitmproxy.proxy.layer import Layer12 13 14class TunnelState(Enum):15    INACTIVE = auto()16    ESTABLISHING = auto()17    OPEN = auto()18    CLOSED = auto()19 20 21class TunnelLayer(layer.Layer):22    """23    A specialized layer that simplifies the implementation of tunneling protocols such as SOCKS, upstream HTTP proxies,24    or TLS.25    """26 27    child_layer: layer.Layer28    tunnel_connection: connection.Connection29    """The 'outer' connection which provides the tunnel protocol I/O"""30    conn: connection.Connection31    """The 'inner' connection which provides data I/O"""32    tunnel_state: TunnelState = TunnelState.INACTIVE33    command_to_reply_to: commands.OpenConnection | None = None34    _event_queue: list[events.Event]35    """36    If the connection already exists when we receive the start event,37    we buffer commands until we have established the tunnel.38    """39 40    def __init__(41        self,42        context: context.Context,43        tunnel_connection: connection.Connection,44        conn: connection.Connection,45    ):46        super().__init__(context)47        self.tunnel_connection = tunnel_connection48        self.conn = conn49        self.child_layer = layer.NextLayer(self.context)50        self._event_queue = []51 52    def __repr__(self):53        return f"{type(self).__name__}({self.tunnel_state.name.lower()})"54 55    def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]:56        if isinstance(event, events.Start):57            if self.tunnel_connection.state is not connection.ConnectionState.CLOSED:58                # we might be in the interesting state here where the connection is already half-closed,59                # for example because next_layer buffered events and the client disconnected in the meantime.60                # we still expect a close event to arrive, so we carry on here as normal for now.61                self.tunnel_state = TunnelState.ESTABLISHING62                yield from self.start_handshake()63            yield from self.event_to_child(event)64        elif (65            isinstance(event, events.ConnectionEvent)66            and event.connection == self.tunnel_connection67        ):68            if isinstance(event, events.DataReceived):69                if self.tunnel_state is TunnelState.ESTABLISHING:70                    done, err = yield from self.receive_handshake_data(event.data)71                    if done:72                        if self.conn != self.tunnel_connection:73                            self.conn.state = connection.ConnectionState.OPEN74                            self.conn.timestamp_start = time.time()75                    if err:76                        if self.conn != self.tunnel_connection:77                            self.conn.state = connection.ConnectionState.CLOSED78                            self.conn.timestamp_start = time.time()79                        yield from self.on_handshake_error(err)80                    if done or err:81                        yield from self._handshake_finished(err)82                else:83                    yield from self.receive_data(event.data)84            elif isinstance(event, events.ConnectionClosed):85                if self.conn != self.tunnel_connection:86                    self.conn.state &= ~connection.ConnectionState.CAN_READ87                    self.conn.timestamp_end = time.time()88                if self.tunnel_state is TunnelState.OPEN:89                    yield from self.receive_close()90                elif self.tunnel_state is TunnelState.ESTABLISHING:91                    err = "connection closed"92                    yield from self.on_handshake_error(err)93                    yield from self._handshake_finished(err)94                self.tunnel_state = TunnelState.CLOSED95            else:  # pragma: no cover96                raise AssertionError(f"Unexpected event: {event}")97        else:98            yield from self.event_to_child(event)99 100    def _handshake_finished(self, err: str | None) -> layer.CommandGenerator[None]:101        if err:102            self.tunnel_state = TunnelState.CLOSED103        else:104            self.tunnel_state = TunnelState.OPEN105        if self.command_to_reply_to:106            yield from self.event_to_child(107                events.OpenConnectionCompleted(self.command_to_reply_to, err)108            )109            self.command_to_reply_to = None110        else:111            for evt in self._event_queue:112                yield from self.event_to_child(evt)113            self._event_queue.clear()114 115    def _handle_command(116        self, command: commands.Command117    ) -> layer.CommandGenerator[None]:118        if (119            isinstance(command, commands.ConnectionCommand)120            and command.connection == self.conn121        ):122            if isinstance(command, commands.SendData):123                yield from self.send_data(command.data)124            elif isinstance(command, commands.CloseConnection):125                if self.conn != self.tunnel_connection:126                    self.conn.state &= ~connection.ConnectionState.CAN_WRITE127                    command.connection = self.tunnel_connection128                yield from self.send_close(command)129            elif isinstance(command, commands.OpenConnection):130                # create our own OpenConnection command object that blocks here.131                self.command_to_reply_to = command132                self.tunnel_state = TunnelState.ESTABLISHING133                err = yield commands.OpenConnection(self.tunnel_connection)134                if err:135                    yield from self.event_to_child(136                        events.OpenConnectionCompleted(command, err)137                    )138                    self.tunnel_state = TunnelState.CLOSED139                else:140                    yield from self.start_handshake()141            else:  # pragma: no cover142                raise AssertionError(f"Unexpected command: {command}")143        else:144            yield command145 146    def event_to_child(self, event: events.Event) -> layer.CommandGenerator[None]:147        if (148            self.tunnel_state is TunnelState.ESTABLISHING149            and not self.command_to_reply_to150        ):151            self._event_queue.append(event)152            return153        for command in self.child_layer.handle_event(event):154            yield from self._handle_command(command)155 156    def start_handshake(self) -> layer.CommandGenerator[None]:157        yield from self._handle_event(events.DataReceived(self.tunnel_connection, b""))158 159    def receive_handshake_data(160        self, data: bytes161    ) -> layer.CommandGenerator[tuple[bool, str | None]]:162        """returns a (done, err) tuple"""163        yield from ()164        return True, None165 166    def on_handshake_error(self, err: str) -> layer.CommandGenerator[None]:167        """Called if either receive_handshake_data returns an error or we receive a close during handshake."""168        yield commands.CloseConnection(self.tunnel_connection)169 170    def receive_data(self, data: bytes) -> layer.CommandGenerator[None]:171        yield from self.event_to_child(events.DataReceived(self.conn, data))172 173    def receive_close(self) -> layer.CommandGenerator[None]:174        yield from self.event_to_child(events.ConnectionClosed(self.conn))175 176    def send_data(self, data: bytes) -> layer.CommandGenerator[None]:177        yield commands.SendData(self.tunnel_connection, data)178 179    def send_close(180        self, command: commands.CloseConnection181    ) -> layer.CommandGenerator[None]:182        yield command183 184 185class LayerStack:186    def __init__(self) -> None:187        self._stack: list[Layer] = []188 189    def __getitem__(self, item: int) -> Layer:190        return self._stack.__getitem__(item)191 192    def __truediv__(self, other: Union[Layer, "LayerStack"]) -> "LayerStack":193        if isinstance(other, Layer):194            if self._stack:195                self._stack[-1].child_layer = other  # type: ignore196            self._stack.append(other)197        else:198            if self._stack:199                self._stack[-1].child_layer = other[0]  # type: ignore200            self._stack.extend(other._stack)201        return self202