Repository: mitmproxy/mitmproxy

Issue:
Support IPv6 transparent proxying on Windows
Transparent mode used to work OK on Windows in version 2.0.2. There was an issue with showing correct Host headers however, so I tried to upgrade to 3.0.4 and found that transparent mode now doesn't work at all. On every client request mitmweb shows the following error message:

> Transparent mode failure: RuntimeError('Cannot resolve original destination.',)

I've tried previous mitmproxy releases and found that this functionality was broken starting from v3.0.0 
RC2. I do understand that this mode is not "officially" supported on Windows, but since it used to work, I hope it can be fixed :)

##### Steps to reproduce the problem:

1. Enable IP forwarding (https://gist.github.com/mhils/e9769fab7b1dd0775335)
2. Set port redirection, e.g.:
> netsh interface portproxy add v4tov4 listenport=80 connectaddress=192.168.0.90 connectport=8080
> netsh interface portproxy add v4tov4 listenport=443 connectaddress=192.168.0.90 connectport=8080
3. Start mitmproxy in transparent mode:
> mitmweb.exe --mode transparent --showhost
4. Install mitmproxy certificate on the phone
5. Set default gateway in phone's Wi-Fi connection properties to the mitmproxy host (e.g. 192.168.0.90)
6. Open any website in phone's browser

##### System information

Mitmproxy: 3.0.4 binary
Python:    3.6.4
OpenSSL:   OpenSSL 1.1.0h  27 Mar 2018
Platform:  Windows-10-10.0.16299



Discussion:
Hi,

2.0.2 still works for you but 3.x doesn't? Also where do the netsh rules come from? Normally that's not required, the important part is to have the redirector running (see https://github.com/mitmproxy/mitmproxy/blob/master/mitmproxy/platform/windows.py#L122)
More specifically, it looks like something was broken between the following commits: https://github.com/mitmproxy/mitmproxy/compare/14b33dca5d19afc2427499e25aee8b7439b4b485...40943f5618d0cc9e5ce100df4620aae9cdc236ac
Since when I'm trying to install specific revision of mitmproxy from source with --hosts issue fixed (commit 40943f5618d0cc9e5ce100df4620aae9cdc236ac), I get the same error (which I don't get in 2.0.2 release).
Can you git-bisect it down to a specific commit?
OK, I'll try to do it later. Regarding your questions above - yes, exactly, 2.0.2 still works but 3.x doesn't. netsh rules come from some mitmproxy setup guide for transparent mode (I just changed Linux iptables port redirection rules to its Windows analogues). When I remove these rules, nothing changes - I see the same error. So probably these rules are no longer necessary, but they are not the cause of the problem anyway.
Not sure why you closed this. The problem is still present in v4 and v3 with only v2 still working for transparent proxying.

Version information I tested this with:
```
Mitmproxy: 4.0.1
Python:    3.6.4
OpenSSL:   OpenSSL 1.1.0h  27 Mar 2018
Platform:  Windows-10-10.0.17134-SP0
```
```
Mitmproxy: 3.0.4
Python:    3.6.4
OpenSSL:   OpenSSL 1.1.0h  27 Mar 2018
Platform:  Windows-10-10.0.17134-SP0
```
```
Mitmproxy version: 2.0.2 (release version)
Python version: 3.6.4
Platform: Windows-10-10.0.17134-SP0
SSL version: OpenSSL 1.1.0f  25 May 2017
Windows version: 10 10.0.17134 SP0 Multiprocessor Free
```
@Jazzelhawk Thank you - I've re-opened this, and we'll investigate further. 
Did a git-bisect to find the specific commit, hope this helps.

```
fab3a8dcf4d7cdce55099172e45a6a0978eac4ab is the first bad commit
commit fab3a8dcf4d7cdce55099172e45a6a0978eac4ab
Author: Maximilian Hils <git@maximilianhils.com>
Date:   Wed Feb 22 22:58:52 2017 +0100

    fix constant definition

:040000 040000 c11646312e07113d064ff98dd56dbb91d6a38fec a0a6c773bd2dcc2a47aefe980c091320a1e65329 M  mitmproxy
```

Resolving patch:
diff --git a/mitmproxy/platform/windows.py b/mitmproxy/platform/windows.py
--- a/mitmproxy/platform/windows.py
+++ b/mitmproxy/platform/windows.py
@@ -1,25 +1,44 @@
-import collections
+import contextlib
 import ctypes
 import ctypes.wintypes
+import io
+import json
 import os
+import re
 import socket
-import struct
+import socketserver
 import threading
 import time
+import typing
 
-import argparse
+import click
+import collections
 import pydivert
 import pydivert.consts
-import pickle
-import socketserver
 
-PROXY_API_PORT = 8085
+REDIRECT_API_HOST = "127.0.0.1"
+REDIRECT_API_PORT = 8085
+
+
+##########################
+# Resolver
+
+def read(rfile: io.BufferedReader) -> typing.Any:
+    x = rfile.readline().strip()
+    return json.loads(x)
+
+
+def write(data, wfile: io.BufferedWriter) -> None:
+    wfile.write(json.dumps(data).encode() + b"\n")
+    wfile.flush()
 
 
 class Resolver:
+    sock: socket.socket
+    lock: threading.RLock
 
     def __init__(self):
-        self.socket = None
+        self.sock = None
         self.lock = threading.RLock()
 
     def setup(self):
@@ -28,406 +47,528 @@ def setup(self):
             self._connect()
 
     def _connect(self):
-        if self.socket:
-            self.socket.close()
-        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self.socket.connect(("127.0.0.1", PROXY_API_PORT))
-
-        self.wfile = self.socket.makefile('wb')
-        self.rfile = self.socket.makefile('rb')
-        pickle.dump(os.getpid(), self.wfile)
-
-    def original_addr(self, csock):
-        client = csock.getpeername()[:2]
+        if self.sock:
+            self.sock.close()
+        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.sock.connect((REDIRECT_API_HOST, REDIRECT_API_PORT))
+
+        self.wfile = self.sock.makefile('wb')
+        self.rfile = self.sock.makefile('rb')
+        write(os.getpid(), self.wfile)
+
+    def original_addr(self, csock: socket.socket):
+        ip, port = csock.getpeername()[:2]
+        ip = re.sub("^::ffff:(?=\d+.\d+.\d+.\d+$)", "", ip)
+        ip = ip.split("%", 1)[0]
         with self.lock:
             try:
-                pickle.dump(client, self.wfile)
-                self.wfile.flush()
-                addr = pickle.load(self.rfile)
+                write((ip, port), self.wfile)
+                addr = read(self.rfile)
                 if addr is None:
                     raise RuntimeError("Cannot resolve original destination.")
-                addr = list(addr)
-                addr[0] = str(addr[0])
-                addr = tuple(addr)
-                return addr
+                return tuple(addr)
             except (EOFError, socket.error):
                 self._connect()
                 return self.original_addr(csock)
 
 
 class APIRequestHandler(socketserver.StreamRequestHandler):
-
     """
     TransparentProxy API: Returns the pickled server address, port tuple
     for each received pickled client address, port tuple.
     """
 
     def handle(self):
-        proxifier = self.server.proxifier
-        pid = None
+        proxifier: TransparentProxy = self.server.proxifier
         try:
-            pid = pickle.load(self.rfile)
-            if pid is not None:
-                proxifier.trusted_pids.add(pid)
-
-            while True:
-                client = pickle.load(self.rfile)
-                server = proxifier.client_server_map.get(client, None)
-                pickle.dump(server, self.wfile)
-                self.wfile.flush()
-
+            pid: int = read(self.rfile)
+            with proxifier.exempt(pid):
+                while True:
+                    client = tuple(read(self.rfile))
+                    try:
+                        server = proxifier.client_server_map[client]
+                    except KeyError:
+                        server = None
+                    write(server, self.wfile)
         except (EOFError, socket.error):
-            proxifier.trusted_pids.discard(pid)
+            pass
 
 
 class APIServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
 
     def __init__(self, proxifier, *args, **kwargs):
-        socketserver.TCPServer.__init__(self, *args, **kwargs)
+        super().__init__(*args, **kwargs)
         self.proxifier = proxifier
         self.daemon_threads = True
 
 
-# Windows error.h
+##########################
+# Windows API
+
+# from Windows' error.h
 ERROR_INSUFFICIENT_BUFFER = 0x7A
 
+IN6_ADDR = ctypes.c_ubyte * 16
+IN4_ADDR = ctypes.c_ubyte * 4
+
+
+#
+# IPv6
+#
+
+# https://msdn.microsoft.com/en-us/library/windows/desktop/aa366896(v=vs.85).aspx
+class MIB_TCP6ROW_OWNER_PID(ctypes.Structure):
+    _fields_ = [
+        ('ucLocalAddr', IN6_ADDR),
+        ('dwLocalScopeId', ctypes.wintypes.DWORD),
+        ('dwLocalPort', ctypes.wintypes.DWORD),
+        ('ucRemoteAddr', IN6_ADDR),
+        ('dwRemoteScopeId', ctypes.wintypes.DWORD),
+        ('dwRemotePort', ctypes.wintypes.DWORD),
+        ('dwState', ctypes.wintypes.DWORD),
+        ('dwOwningPid', ctypes.wintypes.DWORD),
+    ]
+
+
+# https://msdn.microsoft.com/en-us/library/windows/desktop/aa366905(v=vs.85).aspx
+def MIB_TCP6TABLE_OWNER_PID(size):
+    class _MIB_TCP6TABLE_OWNER_PID(ctypes.Structure):
+        _fields_ = [
+            ('dwNumEntries', ctypes.wintypes.DWORD),
+            ('table', MIB_TCP6ROW_OWNER_PID * size)
+        ]
+
+    return _MIB_TCP6TABLE_OWNER_PID()
 
-# http://msdn.microsoft.com/en-us/library/windows/desktop/bb485761(v=vs.85).aspx
-class MIB_TCPROW2(ctypes.Structure):
+
+#
+# IPv4
+#
+
+# https://msdn.microsoft.com/en-us/library/windows/desktop/aa366913(v=vs.85).aspx
+class MIB_TCPROW_OWNER_PID(ctypes.Structure):
     _fields_ = [
         ('dwState', ctypes.wintypes.DWORD),
-        ('dwLocalAddr', ctypes.wintypes.DWORD),
+        ('ucLocalAddr', IN4_ADDR),
         ('dwLocalPort', ctypes.wintypes.DWORD),
-        ('dwRemoteAddr', ctypes.wintypes.DWORD),
+        ('ucRemoteAddr', IN4_ADDR),
         ('dwRemotePort', ctypes.wintypes.DWORD),
         ('dwOwningPid', ctypes.wintypes.DWORD),
-        ('dwOffloadState', ctypes.wintypes.DWORD)
     ]
 
 
-# http://msdn.microsoft.com/en-us/library/windows/desktop/bb485772(v=vs.85).aspx
-def MIB_TCPTABLE2(size):
-    class _MIB_TCPTABLE2(ctypes.Structure):
-        _fields_ = [('dwNumEntries', ctypes.wintypes.DWORD),
-                    ('table', MIB_TCPROW2 * size)]
+# https://msdn.microsoft.com/en-us/library/windows/desktop/aa366921(v=vs.85).aspx
+def MIB_TCPTABLE_OWNER_PID(size):
+    class _MIB_TCPTABLE_OWNER_PID(ctypes.Structure):
+        _fields_ = [
+            ('dwNumEntries', ctypes.wintypes.DWORD),
+            ('table', MIB_TCPROW_OWNER_PID * size)
+        ]
 
-    return _MIB_TCPTABLE2()
+    return _MIB_TCPTABLE_OWNER_PID()
 
 
-class TransparentProxy:
+TCP_TABLE_OWNER_PID_CONNECTIONS = 4
+
+
+class TcpConnectionTable(collections.Mapping):
+    DEFAULT_TABLE_SIZE = 4096
+
+    def __init__(self):
+        self._tcp = MIB_TCPTABLE_OWNER_PID(self.DEFAULT_TABLE_SIZE)
+        self._tcp_size = ctypes.wintypes.DWORD(self.DEFAULT_TABLE_SIZE)
+        self._tcp6 = MIB_TCP6TABLE_OWNER_PID(self.DEFAULT_TABLE_SIZE)
+        self._tcp6_size = ctypes.wintypes.DWORD(self.DEFAULT_TABLE_SIZE)
+        self._map = {}
+
+    def __getitem__(self, item):
+        return self._map[item]
+
+    def __iter__(self):
+        return self._map.__iter__()
+
+    def __len__(self):
+        return self._map.__len__()
+
+    def refresh(self):
+        self._map = {}
+        self._refresh_ipv4()
+        self._refresh_ipv6()
+
+    def _refresh_ipv4(self):
+        ret = ctypes.windll.iphlpapi.GetExtendedTcpTable(
+            ctypes.byref(self._tcp),
+            ctypes.byref(self._tcp_size),
+            False,
+            socket.AF_INET,
+            TCP_TABLE_OWNER_PID_CONNECTIONS,
+            0
+        )
+        if ret == 0:
+            for row in self._tcp.table[:self._tcp.dwNumEntries]:
+                local_ip = socket.inet_ntop(socket.AF_INET, bytes(row.ucLocalAddr))
+                local_port = socket.htons(row.dwLocalPort)
+                self._map[(local_ip, local_port)] = row.dwOwningPid
+        elif ret == ERROR_INSUFFICIENT_BUFFER:
+            self._tcp = MIB_TCPTABLE_OWNER_PID(self._tcp_size.value)
+            # no need to update size, that's already done.
+            self._refresh_ipv4()
+        else:
+            raise RuntimeError("[IPv4] Unknown GetExtendedTcpTable return code: %s" % ret)
+
+    def _refresh_ipv6(self):
+        ret = ctypes.windll.iphlpapi.GetExtendedTcpTable(
+            ctypes.byref(self._tcp6),
+            ctypes.byref(self._tcp6_size),
+            False,
+            socket.AF_INET6,
+            TCP_TABLE_OWNER_PID_CONNECTIONS,
+            0
+        )
+        if ret == 0:
+            for row in self._tcp6.table[:self._tcp6.dwNumEntries]:
+                local_ip = socket.inet_ntop(socket.AF_INET6, bytes(row.ucLocalAddr))
+                local_port = socket.htons(row.dwLocalPort)
+                self._map[(local_ip, local_port)] = row.dwOwningPid
+        elif ret == ERROR_INSUFFICIENT_BUFFER:
+            self._tcp6 = MIB_TCP6TABLE_OWNER_PID(self._tcp6_size.value)
+            # no need to update size, that's already done.
+            self._refresh_ipv6()
+        else:
+            raise RuntimeError("[IPv6] Unknown GetExtendedTcpTable return code: %s" % ret)
+
+
+def get_local_ip() -> typing.Optional[str]:
+    # Auto-Detect local IP. This is required as re-injecting to 127.0.0.1 does not work.
+    # https://stackoverflow.com/questions/166506/finding-local-ip-addresses-using-pythons-stdlib
+    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    try:
+        s.connect(("8.8.8.8", 80))
+        return s.getsockname()[0]
+    except OSError:
+        return None
+    finally:
+        s.close()
+
+
+def get_local_ip6(reachable: str) -> typing.Optional[str]:
+    # The same goes for IPv6, with the added difficulty that .connect() fails if
+    # the target network is not reachable.
+    s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
+    try:
+        s.connect((reachable, 80))
+        return s.getsockname()[0]
+    except OSError:
+        return None
+    finally:
+        s.close()
+
+
+class Redirect(threading.Thread):
+    daemon = True
+    windivert: pydivert.WinDivert
+
+    def __init__(
+        self,
+        handle: typing.Callable[[pydivert.Packet], None],
+        filter: str,
+        layer: pydivert.Layer = pydivert.Layer.NETWORK,
+        flags: pydivert.Flag = 0
+    ) -> None:
+        self.handle = handle
+        self.windivert = pydivert.WinDivert(filter, layer, flags=flags)
+        super().__init__()
+
+    def start(self):
+        self.windivert.open()
+        super().start()
+
+    def run(self):
+        while True:
+            try:
+                packet = self.windivert.recv()
+            except WindowsError as e:
+                if e.winerror == 995:
+                    return
+                else:
+                    raise
+            else:
+                self.handle(packet)
+
+    def shutdown(self):
+        self.windivert.close()
+
+    def recv(self) -> typing.Optional[pydivert.Packet]:
+        """
+        Convenience function that receives a packet from the passed handler and handles error codes.
+        If the process has been shut down, None is returned.
+        """
+        try:
+            return self.windivert.recv()
+        except WindowsError as e:
+            if e.winerror == 995:
+                return None
+            else:
+                raise
+
+
+class RedirectLocal(Redirect):
+    trusted_pids: typing.Set[int]
+
+    def __init__(
+        self,
+        redirect_request: typing.Callable[[pydivert.Packet], None],
+        filter: str
+    ) -> None:
+        self.tcp_connections = TcpConnectionTable()
+        self.trusted_pids = set()
+        self.redirect_request = redirect_request
+        super().__init__(self.handle, filter)
+
+    def handle(self, packet):
+        client = (packet.src_addr, packet.src_port)
+
+        if client not in self.tcp_connections:
+            self.tcp_connections.refresh()
+
+        # If this fails, we most likely have a connection from an external client.
+        # In this, case we always want to proxy the request.
+        pid = self.tcp_connections.get(client, None)
+
+        if pid not in self.trusted_pids:
+            self.redirect_request(packet)
+        else:
+            # It's not really clear why we need to recalculate the checksum here,
+            # but this was identified as necessary in https://github.com/mitmproxy/mitmproxy/pull/3174.
+            self.windivert.send(packet, recalculate_checksum=True)
+
+
+TConnection = typing.Tuple[str, int]
 
+
+class ClientServerMap:
+    """A thread-safe LRU dict."""
+    connection_cache_size: typing.ClassVar[int] = 65536
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._map = collections.OrderedDict()
+
+    def __getitem__(self, item: TConnection) -> TConnection:
+        with self._lock:
+            return self._map[item]
+
+    def __setitem__(self, key: TConnection, value: TConnection) -> None:
+        with self._lock:
+            self._map[key] = value
+            self._map.move_to_end(key)
+            while len(self._map) > self.connection_cache_size:
+                self._map.popitem(False)
+
+
+class TransparentProxy:
     """
-    Transparent Windows Proxy for mitmproxy based on WinDivert/PyDivert.
+    Transparent Windows Proxy for mitmproxy based on WinDivert/PyDivert. This module can be used to
+    redirect both traffic that is forwarded by the host and traffic originating from the host itself.
 
     Requires elevated (admin) privileges. Can be started separately by manually running the file.
 
-    This module can be used to intercept and redirect all traffic that is forwarded by the user's machine and
-    traffic sent from the machine itself.
-
     How it works:
 
-    (1) First, we intercept all packages that match our filter (destination port 80 and 443 by default).
-    We both consider traffic that is forwarded by the OS (WinDivert's NETWORK_FORWARD layer) as well as traffic
-    sent from the local machine (WinDivert's NETWORK layer). In the case of traffic from the local machine, we need to
-    distinguish between traffc sent from applications and traffic sent from the proxy. To accomplish this, we use
-    Windows' GetTcpTable2 syscall to determine the source application's PID.
+    (1) First, we intercept all packages that match our filter.
+    We both consider traffic that is forwarded by the OS (WinDivert's NETWORK_FORWARD layer) as well
+    as traffic sent from the local machine (WinDivert's NETWORK layer). In the case of traffic from
+    the local machine, we need to exempt packets sent from the proxy to not create a redirect loop.
+    To accomplish this, we use Windows' GetExtendedTcpTable syscall and determine the source
+    application's PID.
 
     For each intercepted package, we
         1. Store the source -> destination mapping (address and port)
         2. Remove the package from the network (by not reinjecting it).
-        3. Re-inject the package into the local network stack, but with the destination address changed to the proxy.
+        3. Re-inject the package into the local network stack, but with the destination address
+           changed to the proxy.
 
-    (2) Next, the proxy receives the forwarded packet, but does not know the real destination yet (which we overwrote
-    with the proxy's address). On Linux, we would now call getsockopt(SO_ORIGINAL_DST), but that unfortunately doesn't
-    work on Windows. However, we still have the correct source information. As a workaround, we now access the forward
-    module's API (see APIRequestHandler), submit the source information and get the actual destination back (which the
-    forward module stored in (1.3)).
+    (2) Next, the proxy receives the forwarded packet, but does not know the real destination yet
+    (which we overwrote with the proxy's address). On Linux, we would now call
+    getsockopt(SO_ORIGINAL_DST). We now access the redirect module's API (see APIRequestHandler),
+    submit the source information and get the actual destination back (which we stored in 1.1).
 
-    (3) The proxy now establish the upstream connection as usual.
+    (3) The proxy now establishes the upstream connection as usual.
 
-    (4) Finally, the proxy sends the response back to the client. To make it work, we need to change the packet's source
-    address back to the original destination (using the mapping from (1.3)), to which the client believes he is talking
-    to.
+    (4) Finally, the proxy sends the response back to the client. To make it work, we need to change
+    the packet's source address back to the original destination (using the mapping from 1.1),
+    to which the client believes it is talking to.
 
     Limitations:
 
-    - No IPv6 support. (Pull Requests welcome)
-    - TCP ports do not get re-used simultaneously on the client, i.e. the proxy will fail if application X
-      connects to example.com and example.org from 192.168.0.42:4242 simultaneously. This could be mitigated by
-      introducing unique "meta-addresses" which mitmproxy sees, but this would remove the correct client info from
-      mitmproxy.
-
+    - We assume that ephemeral TCP ports are not re-used for multiple connections at the same time.
+    The proxy will fail if an application connects to example.com and example.org from
+    192.168.0.42:4242 simultaneously. This could be mitigated by introducing unique "meta-addresses"
+    which mitmproxy sees, but this would remove the correct client info from mitmproxy.
     """
+    local: typing.Optional[RedirectLocal] = None
+    # really weird linting error here.
+    forward: typing.Optional[Redirect] = None  # noqa
+    response: Redirect
+    icmp: Redirect
+
+    proxy_port: int
+    filter: str
+
+    client_server_map: ClientServerMap
+
+    def __init__(
+        self,
+        local: bool = True,
+        forward: bool = True,
+        proxy_port: int = 8080,
+        filter: typing.Optional[str] = "tcp.DstPort == 80 or tcp.DstPort == 443",
+    ) -> None:
+        self.proxy_port = proxy_port
+        self.filter = (
+            filter
+            or
+            f"tcp.DstPort != {proxy_port} and tcp.DstPort != {REDIRECT_API_PORT} and tcp.DstPort < 49152"
+        )
 
-    def __init__(self,
-                 mode="both",
-                 redirect_ports=(80, 443), custom_filter=None,
-                 proxy_addr=False, proxy_port=8080,
-                 api_host="localhost", api_port=PROXY_API_PORT,
-                 cache_size=65536):
-        """
-        :param mode: Redirection operation mode: "forward" to only redirect forwarded packets, "local" to only redirect
-        packets originating from the local machine, "both" to redirect both.
-        :param redirect_ports: if the destination port is in this tuple, the requests are redirected to the proxy.
-        :param custom_filter: specify a custom WinDivert filter to select packets that should be intercepted. Overrides
-        redirect_ports setting.
-        :param proxy_addr: IP address of the proxy (IP within a network, 127.0.0.1 does not work). By default,
-        this is detected automatically.
-        :param proxy_port: Port the proxy is listenting on.
-        :param api_host: Host the forward module API is listening on.
-        :param api_port: Port the forward module API is listening on.
-        :param cache_size: Maximum number of connection tuples that are stored. Only relevant in very high
-        load scenarios.
-        """
-        if proxy_port in redirect_ports:
-            raise ValueError("The proxy port must not be a redirect port.")
-
-        if not proxy_addr:
-            # Auto-Detect local IP.
-            # https://stackoverflow.com/questions/166506/finding-local-ip-addresses-using-pythons-stdlib
-            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-            s.connect(("8.8.8.8", 80))
-            proxy_addr = s.getsockname()[0]
-            s.close()
-
-        self.mode = mode
-        self.proxy_addr, self.proxy_port = proxy_addr, proxy_port
-        self.connection_cache_size = cache_size
-
-        self.client_server_map = collections.OrderedDict()
+        self.ipv4_address = get_local_ip()
+        self.ipv6_address = get_local_ip6("2001:4860:4860::8888")
+        # print(f"IPv4: {self.ipv4_address}, IPv6: {self.ipv6_address}")
+        self.client_server_map = ClientServerMap()
 
-        self.api = APIServer(self, (api_host, api_port), APIRequestHandler)
+        self.api = APIServer(self, (REDIRECT_API_HOST, REDIRECT_API_PORT), APIRequestHandler)
         self.api_thread = threading.Thread(target=self.api.serve_forever)
         self.api_thread.daemon = True
 
-        self.request_filter = custom_filter or " or ".join(
-            ("tcp.DstPort == %d" %
-             p) for p in redirect_ports)
-        self.request_forward_handle: pydivert.WinDivert = None
-        self.request_forward_thread = threading.Thread(
-            target=self.request_forward)
-        self.request_forward_thread.daemon = True
-
-        self.addr_pid_map = dict()
-        self.trusted_pids = set()
-        self.tcptable2 = MIB_TCPTABLE2(0)
-        self.tcptable2_size = ctypes.wintypes.DWORD(0)
-        self.request_local_handle: pydivert.WinDivert = None
-        self.request_local_thread = threading.Thread(target=self.request_local)
-        self.request_local_thread.daemon = True
+        if forward:
+            self.forward = Redirect(
+                self.redirect_request,
+                self.filter,
+                pydivert.Layer.NETWORK_FORWARD
+            )
+        if local:
+            self.local = RedirectLocal(
+                self.redirect_request,
+                self.filter
+            )
 
         # The proxy server responds to the client. To the client,
         # this response should look like it has been sent by the real target
-        self.response_filter = "outbound and tcp.SrcPort == %d" % proxy_port
-        self.response_handle: pydivert.WinDivert = None
-        self.response_thread = threading.Thread(target=self.response)
-        self.response_thread.daemon = True
+        self.response = Redirect(
+            self.redirect_response,
+            f"outbound and tcp.SrcPort == {proxy_port}",
+        )
 
-        self.icmp_handle: pydivert.WinDivert = None
+        # Block all ICMP requests (which are sent on Windows by default).
+        # If we don't do this, our proxy machine may send an ICMP redirect to the client,
+        # which instructs the client to directly connect to the real gateway
+        # if they are on the same network.
+        self.icmp = Redirect(
+            lambda _: None,
+            "icmp",
+            flags=pydivert.Flag.DROP
+        )
 
     @classmethod
     def setup(cls):
         # TODO: Make sure that server can be killed cleanly. That's a bit difficult as we don't have access to
         # controller.should_exit when this is called.
         s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        server_unavailable = s.connect_ex(("127.0.0.1", PROXY_API_PORT))
+        server_unavailable = s.connect_ex((REDIRECT_API_HOST, REDIRECT_API_PORT))
         if server_unavailable:
             proxifier = TransparentProxy()
             proxifier.start()
 
     def start(self):
         self.api_thread.start()
-
-        # Block all ICMP requests (which are sent on Windows by default).
-        # In layman's terms: If we don't do this, our proxy machine tells the client that it can directly connect to the
-        # real gateway if they are on the same network.
-        self.icmp_handle = pydivert.WinDivert(
-            filter="icmp",
-            layer=pydivert.Layer.NETWORK,
-            flags=pydivert.Flag.DROP
-        )
-        self.icmp_handle.open()
-
-        self.response_handle = pydivert.WinDivert(
-            filter=self.response_filter,
-            layer=pydivert.Layer.NETWORK
-        )
-        self.response_handle.open()
-        self.response_thread.start()
-
-        if self.mode == "forward" or self.mode == "both":
-            self.request_forward_handle = pydivert.WinDivert(
-                filter=self.request_filter,
-                layer=pydivert.Layer.NETWORK_FORWARD
-            )
-            self.request_forward_handle.open()
-            self.request_forward_thread.start()
-        if self.mode == "local" or self.mode == "both":
-            self.request_local_handle = pydivert.WinDivert(
-                filter=self.request_filter,
-                layer=pydivert.Layer.NETWORK
-            )
-            self.request_local_handle.open()
-            self.request_local_thread.start()
+        self.icmp.start()
+        self.response.start()
+        if self.forward:
+            self.forward.start()
+        if self.local:
+            self.local.start()
 
     def shutdown(self):
-        if self.mode == "local" or self.mode == "both":
-            self.request_local_handle.close()
-        if self.mode == "forward" or self.mode == "both":
-            self.request_forward_handle.close()
-
-        self.response_handle.close()
-        self.icmp_handle.close()
+        if self.local:
+            self.local.shutdown()
+        if self.forward:
+            self.forward.shutdown()
+        self.response.shutdown()
+        self.icmp.shutdown()
         self.api.shutdown()
 
-    def recv(self, handle: pydivert.WinDivert) -> pydivert.Packet:
-        """
-        Convenience function that receives a packet from the passed handler and handles error codes.
-        If the process has been shut down, (None, None) is returned.
-        """
-        try:
-            return handle.recv()
-        except WindowsError as e:
-            if e.winerror == 995:
-                return None
-            else:
-                raise
-
-    def fetch_pids(self):
-        ret = ctypes.windll.iphlpapi.GetTcpTable2(
-            ctypes.byref(
-                self.tcptable2), ctypes.byref(
-                self.tcptable2_size), 0)
-        if ret == ERROR_INSUFFICIENT_BUFFER:
-            self.tcptable2 = MIB_TCPTABLE2(self.tcptable2_size.value)
-            self.fetch_pids()
-        elif ret == 0:
-            for row in self.tcptable2.table[:self.tcptable2.dwNumEntries]:
-                local = (
-                    socket.inet_ntoa(struct.pack('L', row.dwLocalAddr)),
-                    socket.htons(row.dwLocalPort)
-                )
-                self.addr_pid_map[local] = row.dwOwningPid
-        else:
-            raise RuntimeError("Unknown GetTcpTable2 return code: %s" % ret)
-
-    def request_local(self):
-        while True:
-            packet = self.recv(self.request_local_handle)
-            if not packet:
-                return
-
-            client = (packet.src_addr, packet.src_port)
-
-            if client not in self.addr_pid_map:
-                self.fetch_pids()
-
-            # If this fails, we most likely have a connection from an external client to
-            # a local server on 80/443. In this, case we always want to proxy
-            # the request.
-            pid = self.addr_pid_map.get(client, None)
-
-            if pid not in self.trusted_pids:
-                self._request(packet)
-            else:
-                self.request_local_handle.send(packet, recalculate_checksum=False)
-
-    def request_forward(self):
-        """
-        Redirect packages to the proxy
-        """
-        while True:
-            packet = self.recv(self.request_forward_handle)
-            if not packet:
-                return
-
-            self._request(packet)
-
-    def _request(self, packet: pydivert.Packet):
+    def redirect_request(self, packet: pydivert.Packet):
         # print(" * Redirect client -> server to proxy")
-        # print("%s:%s -> %s:%s" % (packet.src_addr, packet.src_port, packet.dst_addr, packet.dst_port))
+        # print(f"{packet.src_addr}:{packet.src_port} -> {packet.dst_addr}:{packet.dst_port}")
         client = (packet.src_addr, packet.src_port)
-        server = (packet.dst_addr, packet.dst_port)
 
-        if client in self.client_server_map:
-            self.client_server_map.move_to_end(client)
+        self.client_server_map[client] = (packet.dst_addr, packet.dst_port)
+
+        # We do need to inject to an external IP here, 127.0.0.1 does not work.
+        if packet.address_family == socket.AF_INET:
+            assert self.ipv4_address
+            packet.dst_addr = self.ipv4_address
+        elif packet.address_family == socket.AF_INET6:
+            if not self.ipv6_address:
+                self.ipv6_address = get_local_ip6(packet.src_addr)
+            assert self.ipv6_address
+            packet.dst_addr = self.ipv6_address
         else:
-            while len(self.client_server_map) > self.connection_cache_size:
-                self.client_server_map.popitem(False)
-            self.client_server_map[client] = server
-
-        packet.dst_addr, packet.dst_port = self.proxy_addr, self.proxy_port
+            raise RuntimeError("Unknown address family")
+        packet.dst_port = self.proxy_port
         packet.direction = pydivert.consts.Direction.INBOUND
 
-        # Use any handle that's on the NETWORK layer - request_local may be
-        # unavailable.
-        self.response_handle.send(packet)
+        # We need a handle on the NETWORK layer. the local handle is not guaranteed to exist,
+        # so we use the response handle.
+        self.response.windivert.send(packet)
 
-    def response(self):
+    def redirect_response(self, packet: pydivert.Packet):
         """
-        Spoof source address of packets send from the proxy to the client
+        If the proxy responds to the client, let the client believe the target server sent the
+        packets.
         """
-        while True:
-            packet = self.recv(self.response_handle)
-            if not packet:
-                return
-
-            # If the proxy responds to the client, let the client believe the target server sent the packets.
-            # print(" * Adjust proxy -> client")
-            client = (packet.dst_addr, packet.dst_port)
-            server = self.client_server_map.get(client, None)
-            if server:
-                packet.src_addr, packet.src_port = server
-                packet.recalculate_checksums()
-            else:
-                print("Warning: Previously unseen connection from proxy to %s:%s." % client)
-
-            self.response_handle.send(packet, recalculate_checksum=False)
+        # print(" * Adjust proxy -> client")
+        client = (packet.dst_addr, packet.dst_port)
+        try:
+            packet.src_addr, packet.src_port = self.client_server_map[client]
+        except KeyError:
+            print(f"Warning: Previously unseen connection from proxy to {client}")
+        else:
+            packet.recalculate_checksums()
 
+        self.response.windivert.send(packet, recalculate_checksum=False)
 
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description="Windows Transparent Proxy"
-    )
-    parser.add_argument(
-        '--mode',
-        choices=[
-            'forward',
-            'local',
-            'both'],
-        default="both",
-        help='redirection operation mode: "forward" to only redirect forwarded packets, '
-        '"local" to only redirect packets originating from the local machine')
-    group = parser.add_mutually_exclusive_group()
-    group.add_argument(
-        "--redirect-ports",
-        nargs="+",
-        type=int,
-        default=[
-            80,
-            443],
-        metavar="80",
-        help="ports that should be forwarded to the proxy")
-    group.add_argument(
-        "--custom-filter",
-        default=None,
-        metavar="WINDIVERT_FILTER",
-        help="Custom WinDivert interception rule.")
-    parser.add_argument("--proxy-addr", default=False,
-                        help="Proxy Server Address")
-    parser.add_argument("--proxy-port", type=int, default=8080,
-                        help="Proxy Server Port")
-    parser.add_argument("--api-host", default="localhost",
-                        help="API hostname to bind to")
-    parser.add_argument("--api-port", type=int, default=PROXY_API_PORT,
-                        help="API port")
-    parser.add_argument("--cache-size", type=int, default=65536,
-                        help="Maximum connection cache size")
-    options = parser.parse_args()
-    proxy = TransparentProxy(**vars(options))
+    @contextlib.contextmanager
+    def exempt(self, pid: int):
+        if self.local:
+            self.local.trusted_pids.add(pid)
+        try:
+            yield
+        finally:
+            if self.local:
+                self.local.trusted_pids.remove(pid)
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.command()
+@click.option("--local/--no-local", default=True,
+              help="Redirect the host's own traffic.")
+@click.option("--forward/--no-forward", default=True,
+              help="Redirect traffic that's forwarded by the host.")
+@click.option("--filter", type=str, metavar="WINDIVERT_FILTER",
+              help="Custom WinDivert interception rule.")
+@click.option("-p", "--proxy-port", type=int, metavar="8080", default=8080,
+              help="The port mitmproxy is listening on.")
+def redirect(**options):
+    """Redirect flows to mitmproxy."""
+    proxy = TransparentProxy(**options)
     proxy.start()
-    print(" * Transparent proxy active.")
-    print("   Filter: {0}".format(proxy.request_filter))
+    print(f" * Redirection active.")
+    print(f"   Filter: {proxy.request_filter}")
     try:
         while True:
             time.sleep(1)
@@ -435,3 +576,16 @@ def response(self):
         print(" * Shutting down...")
         proxy.shutdown()
         print(" * Shut down.")
+
+
+@cli.command()
+def connections():
+    """List all TCP connections and the associated PIDs."""
+    connections = TcpConnectionTable()
+    connections.refresh()
+    for (ip, port), pid in connections.items():
+        print(f"{ip}:{port} -> {pid}")
+
+
+if __name__ == "__main__":
+    cli()


Test patch: