Upstream proxy auth to addon

mitmproxy/addons/__init__.py

@@ -10,6 +10,7 @@ from mitmproxy.addons import serverplayback
 from mitmproxy.addons import stickyauth
 from mitmproxy.addons import stickycookie
 from mitmproxy.addons import streambodies
+from mitmproxy.addons import upstream_proxy_auth
 
 
 def default_addons():
@@ -26,4 +27,5 @@ def default_addons():
         setheaders.SetHeaders(),
         serverplayback.ServerPlayback(),
         clientplayback.ClientPlayback(),
+        upstream_proxy_auth.UpstreamProxyAuth(),
     ]

mitmproxy/addons/upstream_proxy_auth.py

@@ -0,0 +1,30 @@
+import re
+import base64
+
+from mitmproxy import exceptions
+from mitmproxy.utils import strutils
+
+
+def parse_upstream_auth(auth):
+    pattern = re.compile(".+:")
+    if pattern.search(auth) is None:
+        raise exceptions.OptionsError(
+            "Invalid upstream auth specification: %s" % auth
+        )
+    return b"Basic" + b" " + base64.b64encode(strutils.always_bytes(auth))
+
+
+class UpstreamProxyAuth():
+    def __init__(self):
+        self.auth = None
+
+    def configure(self, options, updated):
+        if "upstream_auth" in updated:
+            if options.upstream_auth is None:
+                self.auth = None
+            else:
+                self.auth = parse_upstream_auth(options.upstream_auth)
+
+    def requestheaders(self, f):
+        if self.auth and f.mode == "upstream":
+            f.request.headers["Proxy-Authorization"] = self.auth

mitmproxy/proxy/config.py

@@ -1,11 +1,8 @@
-import base64
 import collections
 import os
 import re
 from typing import Any
 
-from mitmproxy.utils import strutils
-
 from OpenSSL import SSL, crypto
 
 from mitmproxy import exceptions
@@ -56,15 +53,6 @@ def parse_server_spec(spec):
     return ServerSpec(scheme, address)
 
 
-def parse_upstream_auth(auth):
-    pattern = re.compile(".+:")
-    if pattern.search(auth) is None:
-        raise exceptions.OptionsError(
-            "Invalid upstream auth specification: %s" % auth
-        )
-    return b"Basic" + b" " + base64.b64encode(strutils.always_bytes(auth))
-
-
 class ProxyConfig:
 
     def __init__(self, options: moptions.Options) -> None:
@@ -134,11 +122,8 @@ class ProxyConfig:
                 )
 
         self.upstream_server = None
-        self.upstream_auth = None
         if options.upstream_server:
             self.upstream_server = parse_server_spec(options.upstream_server)
-        if options.upstream_auth:
-            self.upstream_auth = parse_upstream_auth(options.upstream_auth)
 
         self.authenticator = authentication.NullProxyAuth(None)
         needsauth = any(

mitmproxy/proxy/protocol/http.py

@@ -228,10 +228,6 @@ class HttpLayer(base.Layer):
         if self.config.options.mode == "reverse":
             f.request.headers["Host"] = self.config.upstream_server.address.host
 
-        # set upstream auth
-        if self.mode is HTTPMode.upstream and self.config.upstream_auth is not None:
-            f.request.headers["Proxy-Authorization"] = self.config.upstream_auth
-
         # Determine .scheme, .host and .port attributes for inline scripts.
         # For absolute-form requests, they are directly given in the request.
         # For authority-form requests, we only need to determine the request scheme.

test/mitmproxy/addons/test_upstream_proxy_auth.py

@@ -0,0 +1,54 @@
+import base64
+
+from mitmproxy import exceptions
+from mitmproxy.test import taddons
+from mitmproxy.test import tflow
+from mitmproxy.test import tutils
+from mitmproxy.addons import upstream_proxy_auth
+
+
+def test_configure():
+    up = upstream_proxy_auth.UpstreamProxyAuth()
+    with taddons.context() as tctx:
+        tctx.configure(up, upstream_auth="test:test")
+        assert up.auth == b"Basic" + b" " + base64.b64encode(b"test:test")
+
+        tctx.configure(up, upstream_auth="test:")
+        assert up.auth == b"Basic" + b" " + base64.b64encode(b"test:")
+
+        tctx.configure(up, upstream_auth=None)
+        assert not up.auth
+
+        tutils.raises(
+            exceptions.OptionsError,
+            tctx.configure,
+            up,
+            upstream_auth=""
+        )
+        tutils.raises(
+            exceptions.OptionsError,
+            tctx.configure,
+            up,
+            upstream_auth=":"
+        )
+        tutils.raises(
+            exceptions.OptionsError,
+            tctx.configure,
+            up,
+            upstream_auth=":test"
+        )
+
+
+def test_simple():
+    up = upstream_proxy_auth.UpstreamProxyAuth()
+    with taddons.context() as tctx:
+        tctx.configure(up, upstream_auth="foo:bar")
+
+        f = tflow.tflow()
+        f.mode = "upstream"
+        up.requestheaders(f)
+        assert "proxy-authorization" in f.request.headers
+
+        f = tflow.tflow()
+        up.requestheaders(f)
+        assert "proxy-authorization" not in f.request.headers

test/mitmproxy/test_proxy.py

@@ -107,14 +107,10 @@ class TestProcessProxyOptions:
         self.assert_noerr("-T")
 
         self.assert_noerr("-U", "http://localhost")
-        self.assert_err("expected one argument", "-U")
         self.assert_err("Invalid server specification", "-U", "upstream")
 
         self.assert_noerr("--upstream-auth", "test:test")
         self.assert_err("expected one argument", "--upstream-auth")
-        self.assert_err(
-            "Invalid upstream auth specification", "--upstream-auth", "test"
-        )
         self.assert_err("mutually exclusive", "-R", "http://localhost", "-T")
 
     def test_socks_auth(self):

test/mitmproxy/test_proxy_config.py

@@ -1,5 +1,4 @@
 from mitmproxy.test import tutils
-import base64
 from mitmproxy.proxy import config
 
 
@@ -26,23 +25,3 @@ def test_parse_server_spec():
         config.parse_server_spec,
         "http://"
     )
-
-
-def test_parse_upstream_auth():
-    tutils.raises(
-        "Invalid upstream auth specification",
-        config.parse_upstream_auth,
-        ""
-    )
-    tutils.raises(
-        "Invalid upstream auth specification",
-        config.parse_upstream_auth,
-        ":"
-    )
-    tutils.raises(
-        "Invalid upstream auth specification",
-        config.parse_upstream_auth,
-        ":test"
-    )
-    assert config.parse_upstream_auth("test:test") == b"Basic" + b" " + base64.b64encode(b"test:test")
-    assert config.parse_upstream_auth("test:") == b"Basic" + b" " + base64.b64encode(b"test:")