From bc01a146b070ecccc4abb5d9382ac4745c430b3c Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sat, 12 Nov 2016 11:39:16 +1300 Subject: [PATCH] Upstream proxy auth to addon --- mitmproxy/addons/__init__.py | 2 + mitmproxy/addons/upstream_proxy_auth.py | 30 +++++++++++ mitmproxy/proxy/config.py | 15 ------ mitmproxy/proxy/protocol/http.py | 4 -- .../addons/test_upstream_proxy_auth.py | 54 +++++++++++++++++++ test/mitmproxy/test_proxy.py | 4 -- test/mitmproxy/test_proxy_config.py | 21 -------- 7 files changed, 86 insertions(+), 44 deletions(-) create mode 100644 mitmproxy/addons/upstream_proxy_auth.py create mode 100644 test/mitmproxy/addons/test_upstream_proxy_auth.py diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index d2b50c35e..d25c231bb 100644 --- a/mitmproxy/addons/__init__.py +++ b/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(), ] diff --git a/mitmproxy/addons/upstream_proxy_auth.py b/mitmproxy/addons/upstream_proxy_auth.py new file mode 100644 index 000000000..2ee51fcbd --- /dev/null +++ b/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 diff --git a/mitmproxy/proxy/config.py b/mitmproxy/proxy/config.py index 9c414b9c4..e144c1757 100644 --- a/mitmproxy/proxy/config.py +++ b/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( diff --git a/mitmproxy/proxy/protocol/http.py b/mitmproxy/proxy/protocol/http.py index 5f4a9856c..ebe41ac3b 100644 --- a/mitmproxy/proxy/protocol/http.py +++ b/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. diff --git a/test/mitmproxy/addons/test_upstream_proxy_auth.py b/test/mitmproxy/addons/test_upstream_proxy_auth.py new file mode 100644 index 000000000..e9a7f4ef4 --- /dev/null +++ b/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 diff --git a/test/mitmproxy/test_proxy.py b/test/mitmproxy/test_proxy.py index 7cadb6c2f..8847c0884 100644 --- a/test/mitmproxy/test_proxy.py +++ b/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): diff --git a/test/mitmproxy/test_proxy_config.py b/test/mitmproxy/test_proxy_config.py index e012cb5ea..e2c39846c 100644 --- a/test/mitmproxy/test_proxy_config.py +++ b/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:")