Merge pull request #3174 from mhils/windows-ipv6

Add support for IPv6 transparent mode on Windows

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 @@ class Resolver:
             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))
+        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.socket.makefile('wb')
-        self.rfile = self.socket.makefile('rb')
-        pickle.dump(os.getpid(), self.wfile)
+        self.wfile = self.sock.makefile('wb')
+        self.rfile = self.sock.makefile('rb')
+        write(os.getpid(), self.wfile)
 
-    def original_addr(self, csock):
-        client = csock.getpeername()[:2]
+    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)
-
+            pid: int = read(self.rfile)
+            with proxifier.exempt(pid):
                 while True:
-                client = pickle.load(self.rfile)
-                server = proxifier.client_server_map.get(client, None)
-                pickle.dump(server, self.wfile)
-                self.wfile.flush()
-
+                    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
 
-# http://msdn.microsoft.com/en-us/library/windows/desktop/bb485761(v=vs.85).aspx
-class MIB_TCPROW2(ctypes.Structure):
+
+#
+# IPv6
+#
+
+# https://msdn.microsoft.com/en-us/library/windows/desktop/aa366896(v=vs.85).aspx
+class MIB_TCP6ROW_OWNER_PID(ctypes.Structure):
     _fields_ = [
-        ('dwState', ctypes.wintypes.DWORD),
-        ('dwLocalAddr', ctypes.wintypes.DWORD),
+        ('ucLocalAddr', IN6_ADDR),
+        ('dwLocalScopeId', ctypes.wintypes.DWORD),
         ('dwLocalPort', ctypes.wintypes.DWORD),
-        ('dwRemoteAddr', ctypes.wintypes.DWORD),
+        ('ucRemoteAddr', IN6_ADDR),
+        ('dwRemoteScopeId', ctypes.wintypes.DWORD),
         ('dwRemotePort', ctypes.wintypes.DWORD),
+        ('dwState', 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/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_TCPTABLE2()
+    return _MIB_TCP6TABLE_OWNER_PID()
 
 
-class TransparentProxy:
+#
+# IPv4
+#
 
-    """
-    Transparent Windows Proxy for mitmproxy based on WinDivert/PyDivert.
+# 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),
+        ('ucLocalAddr', IN4_ADDR),
+        ('dwLocalPort', ctypes.wintypes.DWORD),
+        ('ucRemoteAddr', IN4_ADDR),
+        ('dwRemotePort', ctypes.wintypes.DWORD),
+        ('dwOwningPid', ctypes.wintypes.DWORD),
+    ]
 
-    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.
+# 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)
+        ]
 
-    How it works:
+    return _MIB_TCPTABLE_OWNER_PID()
 
-    (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.
 
-    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.
+TCP_TABLE_OWNER_PID_CONNECTIONS = 4
 
-    (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)).
 
-    (3) The proxy now establish the upstream connection as usual.
+class TcpConnectionTable(collections.Mapping):
+    DEFAULT_TABLE_SIZE = 4096
 
-    (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.
+    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 = {}
 
-    Limitations:
+    def __getitem__(self, item):
+        return self._map[item]
 
-    - 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.
+    def __iter__(self):
+        return self._map.__iter__()
 
-    """
+    def __len__(self):
+        return self._map.__len__()
 
-    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.")
+    def refresh(self):
+        self._map = {}
+        self._refresh_ipv4()
+        self._refresh_ipv6()
 
-        if not proxy_addr:
-            # Auto-Detect local IP.
+    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))
-            proxy_addr = s.getsockname()[0]
+        return s.getsockname()[0]
+    except OSError:
+        return None
+    finally:
         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()
+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()
 
-        self.api = APIServer(self, (api_host, 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
+class Redirect(threading.Thread):
+    daemon = True
+    windivert: pydivert.WinDivert
 
-        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
-
-        # 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.icmp_handle: pydivert.WinDivert = None
-
-    @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))
-        if server_unavailable:
-            proxifier = TransparentProxy()
-            proxifier.start()
+    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.api_thread.start()
+        self.windivert.open()
+        super().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()
+    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):
-        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.windivert.close()
 
-        self.response_handle.close()
-        self.icmp_handle.close()
-        self.api.shutdown()
-
-    def recv(self, handle: pydivert.WinDivert) -> pydivert.Packet:
+    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, None) is returned.
+        If the process has been shut down, None is returned.
         """
         try:
-            return handle.recv()
+            return self.windivert.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
+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.addr_pid_map:
-                self.fetch_pids()
+        if client not in self.tcp_connections:
+            self.tcp_connections.refresh()
 
-            # 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 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._request(packet)
+            self.redirect_request(packet)
         else:
-                self.request_local_handle.send(packet, recalculate_checksum=False)
+            # 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)
 
-    def request_forward(self):
+
+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:
     """
-        Redirect packages to the proxy
+    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.
+
+    How it works:
+
+    (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.
+
+    (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 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.1),
+    to which the client believes it is talking to.
+
+    Limitations:
+
+    - 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.
     """
-        while True:
-            packet = self.recv(self.request_forward_handle)
-            if not packet:
-                return
+    local: typing.Optional[RedirectLocal] = None
+    # really weird linting error here.
+    forward: typing.Optional[Redirect] = None  # noqa
+    response: Redirect
+    icmp: Redirect
 
-            self._request(packet)
+    proxy_port: int
+    filter: str
 
-    def _request(self, packet: pydivert.Packet):
+    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"
+        )
+
+        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, (REDIRECT_API_HOST, REDIRECT_API_PORT), APIRequestHandler)
+        self.api_thread = threading.Thread(target=self.api.serve_forever)
+        self.api_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 = Redirect(
+            self.redirect_response,
+            f"outbound and tcp.SrcPort == {proxy_port}",
+        )
+
+        # 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((REDIRECT_API_HOST, REDIRECT_API_PORT))
+        if server_unavailable:
+            proxifier = TransparentProxy()
+            proxifier.start()
+
+    def start(self):
+        self.api_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.local:
+            self.local.shutdown()
+        if self.forward:
+            self.forward.shutdown()
+        self.response.shutdown()
+        self.icmp.shutdown()
+        self.api.shutdown()
+
+    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()
+        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:
-                print("Warning: Previously unseen connection from proxy to %s:%s." % client)
+            packet.recalculate_checksums()
 
-            self.response_handle.send(packet, recalculate_checksum=False)
+        self.response.windivert.send(packet, recalculate_checksum=False)
+
+    @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)
 
 
-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",
+@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.")
-    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))
+@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 @@ if __name__ == "__main__":
         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()