Add test for upstream-auth

4 files changed, 23 insertions, 3 deletions

diff --git a/mitmproxy/cmdline.py b/mitmproxy/cmdline.py
index e6f0bc20..eb15a1d2 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:
@@ -373,6 +383,7 @@ def proxy_options(parser):
     parser.add_argument(
         "--upstream-auth",
         action="store", dest="upstream_auth", default=None,
+        type=parse_upstream_auth,
         help="""
             Proxy Authentication:
             username:password
diff --git a/mitmproxy/models/http.py b/mitmproxy/models/http.py
index 5d79306b..0338945b 100644
--- a/mitmproxy/models/http.py
+++ b/mitmproxy/models/http.py
@@ -2,7 +2,6 @@ from __future__ import (absolute_import, print_function, division)
 import Cookie
 import copy
 import warnings
-import base64
 from email.utils import parsedate_tz, formatdate, mktime_tz
 import time
 
@@ -194,8 +193,7 @@ class HTTPRequest(MessageMixin, Request):
         return id(self)
 
     def set_auth(self, auth):
-        auth_str = "Basic" + " " + base64.b64encode(auth)
-        self.data.headers.set_all("Proxy-Authorization", (auth_str,))
+        self.data.headers.set_all("Proxy-Authorization", (auth,))
 
     def replace(self, pattern, repl, *args, **kwargs):
         """
diff --git a/test/mitmproxy/test_cmdline.py b/test/mitmproxy/test_cmdline.py
index 5a70f3e0..1546bb41 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,12 @@ def test_parse_server_spec():
         "http://")
 
 
+def test_parse_upstream_auth():
+    tutils.raises("Invalid upstream auth specification", cmdline.parse_upstream_auth, "")
+    assert cmdline.parse_upstream_auth(
+        "test:test") == "Basic" + " " + base64.b64encode("test: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):