Merge pull request #988 from xhy940801/master

Add upstream proxy authentication

6 files changed, 46 insertions, 1 deletions

diff --git a/mitmproxy/cmdline.py b/mitmproxy/cmdline.py
index 3e9fa011..b1b860f8 100644
--- a/mitmproxy/cmdline.py
+++ b/mitmproxy/cmdline.py
@@ -1,6 +1,7 @@
 from __future__ import absolute_import
 import os
 import re
+import base64
 
 import configargparse
 
@@ -117,6 +118,15 @@ def parse_server_spec(url):
     return config.ServerSpec(scheme, address)
 
 
+def parse_upstream_auth(auth):
+    pattern = re.compile(".+:")
+    if pattern.search(auth) is None:
+        raise configargparse.ArgumentTypeError(
+            "Invalid upstream auth specification: %s" % auth
+        )
+    return "Basic" + " " + base64.b64encode(auth)
+
+
 def get_common_options(options):
     stickycookie, stickyauth = None, None
     if options.stickycookie_filt:
@@ -370,6 +380,15 @@ def proxy_options(parser):
             If your OpenSSL version supports ALPN, HTTP/2 is enabled by default.
         """
     )
+    parser.add_argument(
+        "--upstream-auth",
+        action="store", dest="upstream_auth", default=None,
+        type=parse_upstream_auth,
+        help="""
+            Proxy Authentication:
+            username:password
+        """
+    )
     rawtcp = group.add_mutually_exclusive_group()
     rawtcp.add_argument("--raw-tcp", action="store_true", dest="rawtcp")
     rawtcp.add_argument("--no-raw-tcp", action="store_false", dest="rawtcp",
diff --git a/mitmproxy/models/http.py b/mitmproxy/models/http.py
index 394fe51a..0338945b 100644
--- a/mitmproxy/models/http.py
+++ b/mitmproxy/models/http.py
@@ -192,6 +192,9 @@ class HTTPRequest(MessageMixin, Request):
     def __hash__(self):
         return id(self)
 
+    def set_auth(self, auth):
+        self.data.headers.set_all("Proxy-Authorization", (auth,))
+
     def replace(self, pattern, repl, *args, **kwargs):
         """
             Replaces a regular expression pattern with repl in the headers, the
diff --git a/mitmproxy/protocol/http.py b/mitmproxy/protocol/http.py
index 13d7903b..81e59fbb 100644
--- a/mitmproxy/protocol/http.py
+++ b/mitmproxy/protocol/http.py
@@ -179,6 +179,9 @@ class HttpLayer(Layer):
             try:
                 flow = HTTPFlow(self.client_conn, self.server_conn, live=self)
                 flow.request = request
+                # set upstream auth
+                if self.mode == "upstream" and self.config.upstream_auth is not None:
+                    flow.request.set_auth(self.config.upstream_auth)
                 self.process_request_hook(flow)
 
                 if not flow.response:
diff --git a/mitmproxy/proxy/config.py b/mitmproxy/proxy/config.py
index 490cf20c..149d4710 100644
--- a/mitmproxy/proxy/config.py
+++ b/mitmproxy/proxy/config.py
@@ -53,6 +53,7 @@ class ProxyConfig:
             body_size_limit=None,
             mode="regular",
             upstream_server=None,
+            upstream_auth = None,
             authenticator=None,
             ignore_hosts=tuple(),
             tcp_hosts=tuple(),
@@ -77,8 +78,10 @@ class ProxyConfig:
         self.mode = mode
         if upstream_server:
             self.upstream_server = ServerSpec(upstream_server[0], Address.wrap(upstream_server[1]))
+            self.upstream_auth = upstream_auth
         else:
             self.upstream_server = None
+            self.upstream_auth = None
 
         self.check_ignore = HostMatcher(ignore_hosts)
         self.check_tcp = HostMatcher(tcp_hosts)
@@ -110,7 +113,7 @@ def process_proxy_options(parser, options):
     body_size_limit = utils.parse_size(options.body_size_limit)
 
     c = 0
-    mode, upstream_server = "regular", None
+    mode, upstream_server, upstream_auth  = "regular", None, None
     if options.transparent_proxy:
         c += 1
         if not platform.resolver:
@@ -127,6 +130,7 @@ def process_proxy_options(parser, options):
         c += 1
         mode = "upstream"
         upstream_server = options.upstream_proxy
+        upstream_auth = options.upstream_auth
     if c > 1:
         return parser.error(
             "Transparent, SOCKS5, reverse and upstream proxy mode "
@@ -189,6 +193,7 @@ def process_proxy_options(parser, options):
         body_size_limit=body_size_limit,
         mode=mode,
         upstream_server=upstream_server,
+        upstream_auth=upstream_auth,
         ignore_hosts=options.ignore_hosts,
         tcp_hosts=options.tcp_hosts,
         http2=options.http2,
diff --git a/test/mitmproxy/test_cmdline.py b/test/mitmproxy/test_cmdline.py
index 5a70f3e0..e75b7baf 100644
--- a/test/mitmproxy/test_cmdline.py
+++ b/test/mitmproxy/test_cmdline.py
@@ -1,4 +1,5 @@
 import argparse
+import base64
 from mitmproxy import cmdline
 from . import tutils
 
@@ -53,6 +54,16 @@ def test_parse_server_spec():
         "http://")
 
 
+def test_parse_upstream_auth():
+    tutils.raises("Invalid upstream auth specification", cmdline.parse_upstream_auth, "")
+    tutils.raises("Invalid upstream auth specification", cmdline.parse_upstream_auth, ":")
+    tutils.raises("Invalid upstream auth specification", cmdline.parse_upstream_auth, ":test")
+    assert cmdline.parse_upstream_auth(
+        "test:test") == "Basic" + " " + base64.b64encode("test:test")
+    assert cmdline.parse_upstream_auth(
+        "test:") == "Basic" + " " + base64.b64encode("test:")
+
+
 def test_parse_setheaders():
     x = cmdline.parse_setheader("/foo/bar/voing")
     assert x == ("foo", "bar", "voing")
diff --git a/test/mitmproxy/test_proxy.py b/test/mitmproxy/test_proxy.py
index 34b75b62..fddb851e 100644
--- a/test/mitmproxy/test_proxy.py
+++ b/test/mitmproxy/test_proxy.py
@@ -92,6 +92,10 @@ class TestProcessProxyOptions:
         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("not allowed with", "-R", "http://localhost", "-T")
 
     def test_socks_auth(self):