codekingpro/portable-devtools

1from __future__ import annotations2 3import socket4import struct5import sys6from abc import ABCMeta7from collections.abc import Callable8from dataclasses import dataclass9 10from mitmproxy import connection11from mitmproxy.proxy import commands12from mitmproxy.proxy import events13from mitmproxy.proxy import layer14from mitmproxy.proxy.commands import StartHook15from mitmproxy.proxy.mode_specs import ReverseMode16from mitmproxy.proxy.utils import expect17 18if sys.version_info < (3, 11):19    from typing_extensions import assert_never20else:21    from typing import assert_never22 23 24class HttpProxy(layer.Layer):25    @expect(events.Start)26    def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]:27        child_layer = layer.NextLayer(self.context)28        self._handle_event = child_layer.handle_event29        yield from child_layer.handle_event(event)30 31 32class HttpUpstreamProxy(layer.Layer):33    @expect(events.Start)34    def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]:35        child_layer = layer.NextLayer(self.context)36        self._handle_event = child_layer.handle_event37        yield from child_layer.handle_event(event)38 39 40class DestinationKnown(layer.Layer, metaclass=ABCMeta):41    """Base layer for layers that gather connection destination info and then delegate."""42 43    child_layer: layer.Layer44 45    def finish_start(self) -> layer.CommandGenerator[str | None]:46        if (47            self.context.options.connection_strategy == "eager"48            and self.context.server.address49            and self.context.server.transport_protocol == "tcp"50        ):51            err = yield commands.OpenConnection(self.context.server)52            if err:53                self._handle_event = self.done  # type: ignore54                return err55 56        self._handle_event = self.child_layer.handle_event  # type: ignore57        yield from self.child_layer.handle_event(events.Start())58        return None59 60    @expect(events.DataReceived, events.ConnectionClosed)61    def done(self, _) -> layer.CommandGenerator[None]:62        yield from ()63 64 65class ReverseProxy(DestinationKnown):66    @expect(events.Start)67    def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]:68        spec = self.context.client.proxy_mode69        assert isinstance(spec, ReverseMode)70        self.context.server.address = spec.address71 72        self.child_layer = layer.NextLayer(self.context)73 74        # For secure protocols, set SNI if keep_host_header is false75        match spec.scheme:76            case "http3" | "quic" | "https" | "tls" | "dtls":77                if not self.context.options.keep_host_header:78                    self.context.server.sni = spec.address[0]79            case "tcp" | "http" | "udp" | "dns":80                pass81            case _:  # pragma: no cover82                assert_never(spec.scheme)83 84        err = yield from self.finish_start()85        if err:86            yield commands.CloseConnection(self.context.client)87 88 89class TransparentProxy(DestinationKnown):90    @expect(events.Start)91    def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]:92        assert self.context.server.address, "No server address set."93        self.child_layer = layer.NextLayer(self.context)94        err = yield from self.finish_start()95        if err:96            yield commands.CloseConnection(self.context.client)97 98 99SOCKS5_VERSION = 0x05100 101SOCKS5_METHOD_NO_AUTHENTICATION_REQUIRED = 0x00102SOCKS5_METHOD_USER_PASSWORD_AUTHENTICATION = 0x02103SOCKS5_METHOD_NO_ACCEPTABLE_METHODS = 0xFF104 105SOCKS5_ATYP_IPV4_ADDRESS = 0x01106SOCKS5_ATYP_DOMAINNAME = 0x03107SOCKS5_ATYP_IPV6_ADDRESS = 0x04108 109SOCKS5_REP_HOST_UNREACHABLE = 0x04110SOCKS5_REP_COMMAND_NOT_SUPPORTED = 0x07111SOCKS5_REP_ADDRESS_TYPE_NOT_SUPPORTED = 0x08112 113 114@dataclass115class Socks5AuthData:116    client_conn: connection.Client117    username: str118    password: str119    valid: bool = False120 121 122@dataclass123class Socks5AuthHook(StartHook):124    """125    Mitmproxy has received username/password SOCKS5 credentials.126 127    This hook decides whether they are valid by setting `data.valid`.128    """129 130    data: Socks5AuthData131 132 133class Socks5Proxy(DestinationKnown):134    buf: bytes = b""135 136    def socks_err(137        self,138        message: str,139        reply_code: int | None = None,140    ) -> layer.CommandGenerator[None]:141        if reply_code is not None:142            yield commands.SendData(143                self.context.client,144                bytes([SOCKS5_VERSION, reply_code])145                + b"\x00\x01\x00\x00\x00\x00\x00\x00",146            )147        yield commands.CloseConnection(self.context.client)148        yield commands.Log(message)149        self._handle_event = self.done150 151    @expect(events.Start, events.DataReceived, events.ConnectionClosed)152    def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]:153        if isinstance(event, events.Start):154            pass155        elif isinstance(event, events.DataReceived):156            self.buf += event.data157            yield from self.state()158        elif isinstance(event, events.ConnectionClosed):159            if self.buf:160                yield commands.Log(161                    f"Client closed connection before completing SOCKS5 handshake: {self.buf!r}"162                )163            yield commands.CloseConnection(event.connection)164        else:165            raise AssertionError(f"Unknown event: {event}")166 167    def state_greet(self) -> layer.CommandGenerator[None]:168        if len(self.buf) < 2:169            return170 171        if self.buf[0] != SOCKS5_VERSION:172            if self.buf[:3].isupper():173                guess = "Probably not a SOCKS request but a regular HTTP request. "174            else:175                guess = ""176            yield from self.socks_err(177                guess + "Invalid SOCKS version. Expected 0x05, got 0x%x" % self.buf[0]178            )179            return180 181        n_methods = self.buf[1]182        if len(self.buf) < 2 + n_methods:183            return184 185        if "proxyauth" in self.context.options and self.context.options.proxyauth:186            method = SOCKS5_METHOD_USER_PASSWORD_AUTHENTICATION187            self.state = self.state_auth188        else:189            method = SOCKS5_METHOD_NO_AUTHENTICATION_REQUIRED190            self.state = self.state_connect191 192        if method not in self.buf[2 : 2 + n_methods]:193            method_str = (194                "user/password"195                if method == SOCKS5_METHOD_USER_PASSWORD_AUTHENTICATION196                else "no"197            )198            yield from self.socks_err(199                f"Client does not support SOCKS5 with {method_str} authentication.",200                SOCKS5_METHOD_NO_ACCEPTABLE_METHODS,201            )202            return203        yield commands.SendData(self.context.client, bytes([SOCKS5_VERSION, method]))204        self.buf = self.buf[2 + n_methods :]205        yield from self.state()206 207    state: Callable[..., layer.CommandGenerator[None]] = state_greet208 209    def state_auth(self) -> layer.CommandGenerator[None]:210        if len(self.buf) < 3:211            return212 213        # Parsing username and password, which is somewhat atrocious214        user_len = self.buf[1]215        if len(self.buf) < 3 + user_len:216            return217        pass_len = self.buf[2 + user_len]218        if len(self.buf) < 3 + user_len + pass_len:219            return220        user = self.buf[2 : (2 + user_len)].decode("utf-8", "backslashreplace")221        password = self.buf[(3 + user_len) : (3 + user_len + pass_len)].decode(222            "utf-8", "backslashreplace"223        )224 225        data = Socks5AuthData(self.context.client, user, password)226        yield Socks5AuthHook(data)227        if not data.valid:228            # The VER field contains the current **version of the subnegotiation**, which is X'01'.229            yield commands.SendData(self.context.client, b"\x01\x01")230            yield from self.socks_err("authentication failed")231            return232 233        yield commands.SendData(self.context.client, b"\x01\x00")234        self.buf = self.buf[3 + user_len + pass_len :]235        self.state = self.state_connect236        yield from self.state()237 238    def state_connect(self) -> layer.CommandGenerator[None]:239        # Parse Connect Request240        if len(self.buf) < 5:241            return242 243        if self.buf[:3] != b"\x05\x01\x00":244            yield from self.socks_err(245                f"Unsupported SOCKS5 request: {self.buf!r}",246                SOCKS5_REP_COMMAND_NOT_SUPPORTED,247            )248            return249 250        # Determine message length251        atyp = self.buf[3]252        message_len: int253        if atyp == SOCKS5_ATYP_IPV4_ADDRESS:254            message_len = 4 + 4 + 2255        elif atyp == SOCKS5_ATYP_IPV6_ADDRESS:256            message_len = 4 + 16 + 2257        elif atyp == SOCKS5_ATYP_DOMAINNAME:258            message_len = 4 + 1 + self.buf[4] + 2259        else:260            yield from self.socks_err(261                f"Unknown address type: {atyp}", SOCKS5_REP_ADDRESS_TYPE_NOT_SUPPORTED262            )263            return264 265        # Do we have enough bytes yet?266        if len(self.buf) < message_len:267            return268 269        # Parse host and port270        msg, self.buf = self.buf[:message_len], self.buf[message_len:]271 272        host: str273        if atyp == SOCKS5_ATYP_IPV4_ADDRESS:274            host = socket.inet_ntop(socket.AF_INET, msg[4:-2])275        elif atyp == SOCKS5_ATYP_IPV6_ADDRESS:276            host = socket.inet_ntop(socket.AF_INET6, msg[4:-2])277        else:278            host_bytes = msg[5:-2]279            host = host_bytes.decode("ascii", "replace")280 281        (port,) = struct.unpack("!H", msg[-2:])282 283        # We now have all we need, let's get going.284        self.context.server.address = (host, port)285        self.child_layer = layer.NextLayer(self.context)286 287        # this already triggers the child layer's Start event,288        # but that's not a problem in practice...289        err = yield from self.finish_start()290        if err:291            yield commands.SendData(292                self.context.client, b"\x05\x04\x00\x01\x00\x00\x00\x00\x00\x00"293            )294            yield commands.CloseConnection(self.context.client)295        else:296            yield commands.SendData(297                self.context.client, b"\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00"298            )299            if self.buf:300                yield from self.child_layer.handle_event(301                    events.DataReceived(self.context.client, self.buf)302                )303                del self.buf304