codekingpro/portable-devtools

1import time2from logging import DEBUG3 4from h11._receivebuffer import ReceiveBuffer5 6from mitmproxy import connection7from mitmproxy import http8from mitmproxy.net.http import http19from mitmproxy.proxy import commands10from mitmproxy.proxy import context11from mitmproxy.proxy import layer12from mitmproxy.proxy import tunnel13from mitmproxy.proxy.layers import tls14from mitmproxy.proxy.layers.http._hooks import HttpConnectUpstreamHook15from mitmproxy.utils import human16 17 18class HttpUpstreamProxy(tunnel.TunnelLayer):19    buf: ReceiveBuffer20    send_connect: bool21    conn: connection.Server22    tunnel_connection: connection.Server23 24    def __init__(25        self, ctx: context.Context, tunnel_conn: connection.Server, send_connect: bool26    ):27        super().__init__(ctx, tunnel_connection=tunnel_conn, conn=ctx.server)28        self.buf = ReceiveBuffer()29        self.send_connect = send_connect30 31    @classmethod32    def make(cls, ctx: context.Context, send_connect: bool) -> tunnel.LayerStack:33        assert ctx.server.via34        scheme, address = ctx.server.via35        assert scheme in ("http", "https")36 37        http_proxy = connection.Server(address=address)38 39        stack = tunnel.LayerStack()40        if scheme == "https":41            http_proxy.alpn_offers = tls.HTTP1_ALPNS42            http_proxy.sni = address[0]43            stack /= tls.ServerTLSLayer(ctx, http_proxy)44        stack /= cls(ctx, http_proxy, send_connect)45 46        return stack47 48    def start_handshake(self) -> layer.CommandGenerator[None]:49        if not self.send_connect:50            return (yield from super().start_handshake())51        assert self.conn.address52        flow = http.HTTPFlow(self.context.client, self.tunnel_connection)53        authority = (54            self.conn.address[0].encode("idna") + f":{self.conn.address[1]}".encode()55        )56        headers = http.Headers()57        if self.context.options.http_connect_send_host_header:58            headers.insert(0, b"Host", authority)59        flow.request = http.Request(60            host=self.conn.address[0],61            port=self.conn.address[1],62            method=b"CONNECT",63            scheme=b"",64            authority=authority,65            path=b"",66            http_version=b"HTTP/1.1",67            headers=headers,68            content=b"",69            trailers=None,70            timestamp_start=time.time(),71            timestamp_end=time.time(),72        )73        yield HttpConnectUpstreamHook(flow)74        raw = http1.assemble_request(flow.request)75        yield commands.SendData(self.tunnel_connection, raw)76 77    def receive_handshake_data(78        self, data: bytes79    ) -> layer.CommandGenerator[tuple[bool, str | None]]:80        if not self.send_connect:81            return (yield from super().receive_handshake_data(data))82        self.buf += data83        response_head = self.buf.maybe_extract_lines()84        if response_head:85            try:86                response = http1.read_response_head([bytes(x) for x in response_head])87            except ValueError as e:88                proxyaddr = human.format_address(self.tunnel_connection.address)89                yield commands.Log(f"{proxyaddr}: {e}")90                return False, f"Error connecting to {proxyaddr}: {e}"91            if 200 <= response.status_code < 300:92                if self.buf:93                    yield from self.receive_data(bytes(self.buf))94                    del self.buf95                return True, None96            else:97                proxyaddr = human.format_address(self.tunnel_connection.address)98                raw_resp = b"\n".join(response_head)99                yield commands.Log(f"{proxyaddr}: {raw_resp!r}", DEBUG)100                return (101                    False,102                    f"Upstream proxy {proxyaddr} refused HTTP CONNECT request: {response.status_code} {response.reason}",103                )104        else:105            return False, None106