codekingpro/portable-devtools

1import base642import re3from typing import Optional4 5from mitmproxy import ctx6from mitmproxy import exceptions7from mitmproxy import http8from mitmproxy.proxy import mode_specs9from mitmproxy.utils import strutils10 11 12def parse_upstream_auth(auth: str) -> bytes:13    pattern = re.compile(".+:")14    if pattern.search(auth) is None:15        raise exceptions.OptionsError("Invalid upstream auth specification: %s" % auth)16    return b"Basic" + b" " + base64.b64encode(strutils.always_bytes(auth))17 18 19class UpstreamAuth:20    """21    This addon handles authentication to systems upstream from us for the22    upstream proxy and reverse proxy mode. There are 3 cases:23 24    - Upstream proxy CONNECT requests should have authentication added, and25      subsequent already connected requests should not.26    - Upstream proxy regular requests27    - Reverse proxy regular requests (CONNECT is invalid in this mode)28    """29 30    auth: bytes | None = None31 32    def load(self, loader):33        loader.add_option(34            "upstream_auth",35            Optional[str],36            None,37            """38            Add HTTP Basic authentication to upstream proxy and reverse proxy39            requests. Format: username:password.40            """,41        )42 43    def configure(self, updated):44        if "upstream_auth" in updated:45            if ctx.options.upstream_auth is None:46                self.auth = None47            else:48                self.auth = parse_upstream_auth(ctx.options.upstream_auth)49 50    def http_connect_upstream(self, f: http.HTTPFlow):51        if self.auth:52            f.request.headers["Proxy-Authorization"] = self.auth53 54    def requestheaders(self, f: http.HTTPFlow):55        if self.auth:56            if (57                isinstance(f.client_conn.proxy_mode, mode_specs.UpstreamMode)58                and f.request.scheme == "http"59            ):60                f.request.headers["Proxy-Authorization"] = self.auth61            elif isinstance(f.client_conn.proxy_mode, mode_specs.ReverseMode):62                f.request.headers["Authorization"] = self.auth63