Fix transparent mode initialisation, refactor mitmproxy.platform (#1787)

mitmproxy/platform/__init__.py

@@ -1,17 +1,35 @@
-import sys
 import re
+import socket
+import sys
+from typing import Tuple
+
+
+def init_transparent_mode() -> None:
+    """
+    Initialize transparent mode.
+    """
+
+
+def original_addr(csock: socket.socket) -> Tuple[str, int]:
+    """
+    Get the original destination for the given socket.
+    This function will be None if transparent mode is not supported.
+    """
 
-resolver = None
 
 if re.match(r"linux(?:2)?", sys.platform):
     from . import linux
-    resolver = linux.Resolver
-elif sys.platform == "darwin":
+
+    original_addr = linux.original_addr  # noqa
+elif sys.platform == "darwin" or sys.platform.startswith("freebsd"):
     from . import osx
-    resolver = osx.Resolver
-elif sys.platform.startswith("freebsd"):
-    from . import osx
-    resolver = osx.Resolver
+
+    original_addr = osx.original_addr  # noqa
 elif sys.platform == "win32":
     from . import windows
-    resolver = windows.Resolver
+
+    resolver = windows.Resolver()
+    init_transparent_mode = resolver.setup  # noqa
+    original_addr = resolver.original_addr  # noqa
+else:
+    original_addr = None  # noqa

mitmproxy/platform/linux.py

@@ -5,10 +5,8 @@ import struct
 SO_ORIGINAL_DST = 80
 
 
-class Resolver:
-
-    def original_addr(self, csock):
-        odestdata = csock.getsockopt(socket.SOL_IP, SO_ORIGINAL_DST, 16)
-        _, port, a1, a2, a3, a4 = struct.unpack("!HHBBBBxxxxxxxx", odestdata)
-        address = "%d.%d.%d.%d" % (a1, a2, a3, a4)
-        return address, port
+def original_addr(csock: socket.socket):
+    odestdata = csock.getsockopt(socket.SOL_IP, SO_ORIGINAL_DST, 16)
+    _, port, a1, a2, a3, a4 = struct.unpack("!HHBBBBxxxxxxxx", odestdata)
+    address = "%d.%d.%d.%d" % (a1, a2, a3, a4)
+    return address, port

mitmproxy/platform/osx.py

@@ -14,24 +14,23 @@ from . import pf
     the output processing of pfctl (see pf.py).
 """
 
+STATECMD = ("sudo", "-n", "/sbin/pfctl", "-s", "state")
 
-class Resolver:
-    STATECMD = ("sudo", "-n", "/sbin/pfctl", "-s", "state")
 
-    def original_addr(self, csock):
-        peer = csock.getpeername()
-        try:
-            stxt = subprocess.check_output(self.STATECMD, stderr=subprocess.STDOUT)
-        except subprocess.CalledProcessError as e:
-            if "sudo: a password is required" in e.output.decode(errors="replace"):
-                insufficient_priv = True
-            else:
-                raise RuntimeError("Error getting pfctl state: " + repr(e))
+def original_addr(csock):
+    peer = csock.getpeername()
+    try:
+        stxt = subprocess.check_output(STATECMD, stderr=subprocess.STDOUT)
+    except subprocess.CalledProcessError as e:
+        if "sudo: a password is required" in e.output.decode(errors="replace"):
+            insufficient_priv = True
         else:
-            insufficient_priv = "sudo: a password is required" in stxt.decode(errors="replace")
+            raise RuntimeError("Error getting pfctl state: " + repr(e))
+    else:
+        insufficient_priv = "sudo: a password is required" in stxt.decode(errors="replace")
 
-        if insufficient_priv:
-            raise RuntimeError(
-                "Insufficient privileges to access pfctl. "
-                "See http://docs.mitmproxy.org/en/latest/transparent/osx.html for details.")
-        return pf.lookup(peer[0], peer[1], stxt)
+    if insufficient_priv:
+        raise RuntimeError(
+            "Insufficient privileges to access pfctl. "
+            "See http://docs.mitmproxy.org/en/latest/transparent/osx.html for details.")
+    return pf.lookup(peer[0], peer[1], stxt)

mitmproxy/platform/windows.py

@@ -19,10 +19,13 @@ PROXY_API_PORT = 8085
 class Resolver:
 
     def __init__(self):
-        TransparentProxy.setup()
         self.socket = None
         self.lock = threading.RLock()
-        self._connect()
+
+    def setup(self):
+        with self.lock:
+            TransparentProxy.setup()
+            self._connect()
 
     def _connect(self):
         if self.socket:

mitmproxy/proxy/modes/transparent_proxy.py

@@ -7,11 +7,10 @@ class TransparentProxy(protocol.Layer, protocol.ServerConnectionMixin):
 
     def __init__(self, ctx):
         super().__init__(ctx)
-        self.resolver = platform.resolver()
 
     def __call__(self):
         try:
-            self.server_conn.address = self.resolver.original_addr(self.client_conn.connection)
+            self.server_conn.address = platform.original_addr(self.client_conn.connection)
         except Exception as e:
             raise exceptions.ProtocolException("Transparent mode failure: %s" % repr(e))
 

mitmproxy/proxy/server.py

@@ -6,6 +6,8 @@ from mitmproxy import exceptions
 from mitmproxy import connections
 from mitmproxy import http
 from mitmproxy import log
+from mitmproxy import platform
+from mitmproxy.proxy import ProxyConfig
 from mitmproxy.proxy import modes
 from mitmproxy.proxy import root_context
 from mitmproxy.net import tcp
@@ -33,7 +35,7 @@ class ProxyServer(tcp.TCPServer):
     allow_reuse_address = True
     bound = True
 
-    def __init__(self, config):
+    def __init__(self, config: ProxyConfig):
         """
             Raises ServerException if there's a startup problem.
         """
@@ -42,6 +44,8 @@ class ProxyServer(tcp.TCPServer):
             super().__init__(
                 (config.options.listen_host, config.options.listen_port)
             )
+            if config.options.mode == "transparent":
+                platform.init_transparent_mode()
         except socket.error as e:
             raise exceptions.ServerException(
                 'Error starting proxy server: ' + repr(e)

mitmproxy/tools/cmdline.py

@@ -174,7 +174,7 @@ def get_common_options(args):
     mode, upstream_server = "regular", None
     if args.transparent_proxy:
         c += 1
-        if not platform.resolver:
+        if not platform.original_addr:
             raise exceptions.OptionsError(
                 "Transparent mode not supported on this platform."
             )

test/mitmproxy/test_proxy.py

@@ -94,11 +94,11 @@ class TestProcessProxyOptions:
         with tutils.tmpdir() as cadir:
             self.assert_noerr("--cadir", cadir)
 
-    @mock.patch("mitmproxy.platform.resolver", None)
+    @mock.patch("mitmproxy.platform.original_addr", None)
     def test_no_transparent(self):
         self.assert_err("transparent mode not supported", "-T")
 
-    @mock.patch("mitmproxy.platform.resolver")
+    @mock.patch("mitmproxy.platform.original_addr")
     def test_modes(self, _):
         self.assert_noerr("-R", "http://localhost")
         self.assert_err("expected one argument", "-R")

test/mitmproxy/test_server.py

@@ -2,6 +2,8 @@ import os
 import socket
 import time
 
+import mock
+
 from mitmproxy.test import tutils
 from mitmproxy import controller
 from mitmproxy import options
@@ -878,16 +880,10 @@ class TestKillResponse(tservers.HTTPProxyTest):
         assert self.server.last_log()
 
 
-class EResolver(tservers.TResolver):
-
-    def original_addr(self, sock):
-        raise RuntimeError("Could not resolve original destination.")
-
-
 class TestTransparentResolveError(tservers.TransparentProxyTest):
-    resolver = EResolver
-
-    def test_resolve_error(self):
+    @mock.patch("mitmproxy.platform.original_addr")
+    def test_resolve_error(self, original_addr):
+        original_addr.side_effect = RuntimeError
         assert self.pathod("304").status_code == 502
 
 

test/mitmproxy/tservers.py

@@ -1,9 +1,9 @@
 import os.path
 import threading
 import tempfile
-import mock
 import sys
 
+import mitmproxy.platform
 from mitmproxy.proxy.config import ProxyConfig
 from mitmproxy.proxy.server import ProxyServer
 from mitmproxy import master
@@ -207,33 +207,22 @@ class HTTPProxyTest(ProxyTestBase):
                 return p.request("get:'http://%s%s'" % (options.APP_HOST, page))
 
 
-class TResolver:
-
-    def __init__(self, port):
-        self.port = port
-
-    def original_addr(self, sock):
-        return ("127.0.0.1", self.port)
-
-
 class TransparentProxyTest(ProxyTestBase):
     ssl = None
-    resolver = TResolver
 
     @classmethod
     def setup_class(cls):
+        cls._init_transparent_mode = mitmproxy.platform.init_transparent_mode
+        cls._original_addr = mitmproxy.platform.original_addr
+        mitmproxy.platform.init_transparent_mode = lambda: True
+        mitmproxy.platform.original_addr = lambda sock: ("127.0.0.1", cls.server.port)
         super().setup_class()
 
-        cls._resolver = mock.patch(
-            "mitmproxy.platform.resolver",
-            new=lambda: cls.resolver(cls.server.port)
-        )
-        cls._resolver.start()
-
     @classmethod
     def teardown_class(cls):
-        cls._resolver.stop()
         super().teardown_class()
+        mitmproxy.platform.init_transparent_mode = cls._init_transparent_mode
+        mitmproxy.platform.original_addr = cls._original_addr
 
     @classmethod
     def get_options(cls):