fix #318

2 files changed, 64 insertions, 46 deletions

diff --git a/libmproxy/proxy/config.py b/libmproxy/proxy/config.py
index 94a12bf4..a6a962ea 100644
--- a/libmproxy/proxy/config.py
+++ b/libmproxy/proxy/config.py
@@ -11,14 +11,32 @@ CONF_DIR = "~/.mitmproxy"
 
 class ProxyConfig:
     def __init__(self, confdir=CONF_DIR, clientcerts=None,
-                       no_upstream_cert=False, body_size_limit=None, get_upstream_server=None,
-                       http_form_in="absolute", http_form_out="relative", authenticator=None,
-                       ciphers=None, certs=[], certforward = False
-                ):
+                 no_upstream_cert=False, body_size_limit=None,
+                 mode=None, upstream_server=None, http_form_in=None, http_form_out=None,
+                 authenticator=None,
+                 ciphers=None, certs=[], certforward=False):
         self.ciphers = ciphers
         self.clientcerts = clientcerts
         self.no_upstream_cert = no_upstream_cert
         self.body_size_limit = body_size_limit
+
+        if mode == "transparent":
+            get_upstream_server = TransparentUpstreamServerResolver(platform.resolver(), TRANSPARENT_SSL_PORTS)
+            http_form_in_default, http_form_out_default = "relative", "relative"
+        elif mode == "reverse":
+            get_upstream_server = ConstUpstreamServerResolver(upstream_server)
+            http_form_in_default, http_form_out_default = "relative", "relative"
+        elif mode == "upstream":
+            get_upstream_server = ConstUpstreamServerResolver(upstream_server)
+            http_form_in_default, http_form_out_default = "absolute", "absolute"
+        elif upstream_server:
+            get_upstream_server = ConstUpstreamServerResolver(upstream_server)
+            http_form_in_default, http_form_out_default = "absolute", "relative"
+        else:
+            get_upstream_server, http_form_in_default, http_form_out_default = None, "absolute", "relative"
+        http_form_in = http_form_in or http_form_in_default
+        http_form_out = http_form_out or http_form_out_default
+
         self.get_upstream_server = get_upstream_server
         self.http_form_in = http_form_in
         self.http_form_out = http_form_out
@@ -35,32 +53,27 @@ def process_proxy_options(parser, options):
     body_size_limit = utils.parse_size(options.body_size_limit)
 
     c = 0
-    http_form_in, http_form_out = "absolute", "relative"
-    get_upstream_server = None
+    mode, upstream_server = None, None
     if options.transparent_proxy:
         c += 1
         if not platform.resolver:
             return parser.error("Transparent mode not supported on this platform.")
-        get_upstream_server = TransparentUpstreamServerResolver(platform.resolver(), TRANSPARENT_SSL_PORTS)
-        http_form_in, http_form_out = "relative", "relative"
+        mode = "transparent"
     if options.reverse_proxy:
         c += 1
-        get_upstream_server = ConstUpstreamServerResolver(options.reverse_proxy)
-        http_form_in, http_form_out = "relative", "relative"
+        mode = "reverse"
+        upstream_server = options.reverse_proxy
     if options.upstream_proxy:
         c += 1
-        get_upstream_server = ConstUpstreamServerResolver(options.upstream_proxy)
-        http_form_in, http_form_out = "absolute", "absolute"
+        mode = "upstream"
+        upstream_server = options.upstream_proxy
     if options.manual_destination_server:
         c += 1
-        get_upstream_server = ConstUpstreamServerResolver(options.manual_destination_server)
+        mode = "manual"
+        upstream_server = options.manual_destination_server
     if c > 1:
         return parser.error("Transparent mode, reverse mode, upstream proxy mode and "
                             "specification of an upstream server are mutually exclusive.")
-    if options.http_form_in:
-        http_form_in = options.http_form_in
-    if options.http_form_out:
-        http_form_out = options.http_form_out
 
     if options.clientcerts:
         options.clientcerts = os.path.expanduser(options.clientcerts)
@@ -93,21 +106,22 @@ def process_proxy_options(parser, options):
             parts = ["*", parts[0]]
         parts[1] = os.path.expanduser(parts[1])
         if not os.path.exists(parts[1]):
-            parser.error("Certificate file does not exist: %s"%parts[1])
+            parser.error("Certificate file does not exist: %s" % parts[1])
         certs.append(parts)
 
     return ProxyConfig(
-        clientcerts = options.clientcerts,
-        body_size_limit = body_size_limit,
-        no_upstream_cert = options.no_upstream_cert,
-        get_upstream_server = get_upstream_server,
-        confdir = options.confdir,
-        http_form_in = http_form_in,
-        http_form_out = http_form_out,
-        authenticator = authenticator,
-        ciphers = options.ciphers,
-        certs = certs,
-        certforward = options.certforward,
+        confdir=options.confdir,
+        clientcerts=options.clientcerts,
+        no_upstream_cert=options.no_upstream_cert,
+        body_size_limit=body_size_limit,
+        mode=mode,
+        upstream_server=upstream_server,
+        http_form_in=options.http_form_in,
+        http_form_out=options.http_form_out,
+        authenticator=authenticator,
+        ciphers=options.ciphers,
+        certs=certs,
+        certforward=options.certforward,
     )
 
 
@@ -115,10 +129,10 @@ def ssl_option_group(parser):
     group = parser.add_argument_group("SSL")
     group.add_argument(
         "--cert", dest='certs', default=[], type=str,
-        metavar = "SPEC", action="append",
-        help='Add an SSL certificate. SPEC is of the form "[domain=]path". '\
-             'The domain may include a wildcard, and is equal to "*" if not specified. '\
-             'The file at path is a certificate in PEM format. If a private key is included in the PEM, '\
+        metavar="SPEC", action="append",
+        help='Add an SSL certificate. SPEC is of the form "[domain=]path". ' \
+             'The domain may include a wildcard, and is equal to "*" if not specified. ' \
+             'The file at path is a certificate in PEM format. If a private key is included in the PEM, ' \
              'it is used, else the default key in the conf dir is used. Can be passed multiple times.'
     )
     group.add_argument(
diff --git a/test/tservers.py b/test/tservers.py
index 146dd13d..597ad4ee 100644
--- a/test/tservers.py
+++ b/test/tservers.py
@@ -2,6 +2,8 @@ import os.path
 import threading, Queue
 import shutil, tempfile
 import flask
+import mock
+
 from libmproxy.proxy.config import ProxyConfig
 from libmproxy.proxy.server import ProxyServer
 from libmproxy.proxy.primitives import TransparentUpstreamServerResolver
@@ -88,28 +90,25 @@ class ProxTestBase(object):
         cls.server2 = libpathod.test.Daemon(ssl=cls.ssl, ssloptions=cls.ssloptions)
         pconf = cls.get_proxy_config()
         cls.confdir = os.path.join(tempfile.gettempdir(), "mitmproxy")
-        config = ProxyConfig(
+        cls.config = ProxyConfig(
             no_upstream_cert = cls.no_upstream_cert,
             confdir = cls.confdir,
             authenticator = cls.authenticator,
             certforward = cls.certforward,
             **pconf
         )
-        tmaster = cls.masterclass(config)
+        tmaster = cls.masterclass(cls.config)
         tmaster.start_app(APP_HOST, APP_PORT, cls.externalapp)
         cls.proxy = ProxyThread(tmaster)
         cls.proxy.start()
 
-    @classmethod
-    def tearDownAll(cls):
-        shutil.rmtree(cls.confdir)
-
     @property
     def master(cls):
         return cls.proxy.tmaster
 
     @classmethod
     def teardownAll(cls):
+        shutil.rmtree(cls.confdir)
         cls.proxy.shutdown()
         cls.server.shutdown()
         cls.server2.shutdown()
@@ -189,16 +188,21 @@ class TResolver:
 class TransparentProxTest(ProxTestBase):
     ssl = None
     resolver = TResolver
+
     @classmethod
-    def get_proxy_config(cls):
-        d = ProxTestBase.get_proxy_config()
+    @mock.patch("libmproxy.platform.resolver")
+    def setupAll(cls, _):
+        super(TransparentProxTest, cls).setupAll()
         if cls.ssl:
             ports = [cls.server.port, cls.server2.port]
         else:
             ports = []
-        d["get_upstream_server"] = TransparentUpstreamServerResolver(cls.resolver(cls.server.port), ports)
-        d["http_form_in"] = "relative"
-        d["http_form_out"] = "relative"
+        cls.config.get_upstream_server = TransparentUpstreamServerResolver(cls.resolver(cls.server.port), ports)
+
+    @classmethod
+    def get_proxy_config(cls):
+        d = ProxTestBase.get_proxy_config()
+        d["mode"] = "transparent"
         return d
 
     def pathod(self, spec, sni=None):
@@ -227,7 +231,7 @@ class ReverseProxTest(ProxTestBase):
     @classmethod
     def get_proxy_config(cls):
         d = ProxTestBase.get_proxy_config()
-        d["get_upstream_server"] = lambda c: (
+        d["upstream_server"] = (
             True if cls.ssl else False,
             True if cls.ssl else False,
             "127.0.0.1",
@@ -264,7 +268,7 @@ class ChainProxTest(ProxTestBase):
     """
     n = 2
     chain_config = [lambda port: ProxyConfig(
-        get_upstream_server = lambda c: (False, False, "127.0.0.1", port),
+        upstream_server= (False, False, "127.0.0.1", port),
         http_form_in = "absolute",
         http_form_out = "absolute"
     )] * n