[sans-io] implement layer selection as an addon

mitmproxy/addons/__init__.py

@@ -9,6 +9,7 @@ from mitmproxy.addons import core
 from mitmproxy.addons import cut
 from mitmproxy.addons import disable_h2c
 from mitmproxy.addons import export
+from mitmproxy.addons import next_layer
 from mitmproxy.addons import onboarding
 from mitmproxy.addons import proxyserver
 from mitmproxy.addons import proxyauth
@@ -38,6 +39,7 @@ def default_addons():
         cut.Cut(),
         disable_h2c.DisableH2C(),
         export.Export(),
+        next_layer.NextLayer(),
         onboarding.Onboarding(),
         proxyauth.ProxyAuth(),
         proxyserver.Proxyserver(),

mitmproxy/addons/next_layer.py

@@ -0,0 +1,41 @@
+from mitmproxy import ctx
+from mitmproxy.net import server_spec
+from mitmproxy.proxy.protocol import is_tls_record_magic
+from mitmproxy.proxy2 import layer, layers
+
+
+class NextLayer:
+    def next_layer(self, nextlayer: layer.NextLayer):
+        top_layer = nextlayer.context.layers[-1]
+        data_client = nextlayer.data_client()
+
+        if len(data_client) < 3:
+            return
+
+        client_tls = is_tls_record_magic(data_client)
+
+        # 1. check for --ignore
+
+        # 2. Always insert a TLS layer as second layer, even if there's neither client nor server
+        # tls. An addon may upgrade from http to https, in which case we need a TLS layer.
+        if isinstance(top_layer, layers.modes.ReverseProxy):
+            if client_tls:
+                nextlayer.layer = layers.TLSLayer(
+                    nextlayer.context,
+                    client_tls,
+                    server_spec.parse_with_mode(ctx.options.mode)[1].scheme == "https"
+                )
+            else:
+                # FIXME: TLSLayer doesn't support non-TLS yet, so remove this here once that's in.
+                nextlayer.layer = layers.HTTPLayer(
+                    nextlayer.context
+                )
+        # TODO: Other top layers
+
+        pass
+        # 3. In Http Proxy mode and Upstream Proxy mode, the next layer is fixed.
+        # 4. Check for other TLS cases (e.g. after CONNECT).
+        # 5. Check for --tcp
+        # 6. Check for TLS ALPN (HTTP1/HTTP2)
+        # 7. Check for raw tcp mode
+        # 8. Assume HTTP1 by default

mitmproxy/addons/proxyserver.py

@@ -52,10 +52,11 @@ class Proxyserver:
         self.loop = asyncio.get_event_loop()
         self.listen_port = None
         self.event_queue = None
-        self.options = ctx.options
+        self.options = None
         self._lock = asyncio.Lock()
 
     def running(self):
+        self.options = ctx.options
         self.event_queue = ctx.master.event_queue
         threading.Thread(target=self.loop.run_forever, daemon=True).start()
 

mitmproxy/proxy2/layer.py

@@ -3,10 +3,10 @@ Base class for protocol layers.
 """
 import collections
 import typing
-from abc import ABCMeta, abstractmethod
+from abc import abstractmethod
 
 from mitmproxy.proxy2 import commands, events
-from mitmproxy.proxy2.context import Context
+from mitmproxy.proxy2.context import Context, Connection
 
 
 class Paused(typing.NamedTuple):
@@ -17,12 +17,13 @@ class Paused(typing.NamedTuple):
     generator: commands.TCommandGenerator
 
 
-class Layer(metaclass=ABCMeta):
+class Layer:
     context: Context
     _paused: typing.Optional[Paused]
 
     def __init__(self, context: Context) -> None:
         self.context = context
+        self.context.layers.append(self)
         self._paused = None
         self._paused_event_queue: typing.Deque[events.Event] = collections.deque()
 
@@ -35,7 +36,7 @@ class Layer(metaclass=ABCMeta):
         if False:
             yield None
 
-    def handle_event(self, event: events. Event) -> commands.TCommandGenerator:
+    def handle_event(self, event: events.Event) -> commands.TCommandGenerator:
         if self._paused:
             # did we just receive the reply we were waiting for?
             pause_finished = (
@@ -89,3 +90,52 @@ class Layer(metaclass=ABCMeta):
             command_generator = self._handle_event(event)
             yield from self.__process(command_generator)
             self._debug("#>")
+
+
+mevents = events  # alias here because autocomplete above should not have aliased version.
+
+
+class NextLayer(Layer):
+    layer: typing.Optional[Layer]
+    """The next layer. To be set by an addon."""
+
+    events: typing.List[mevents.Event]
+    """All events that happened before a decision was made."""
+
+    def __init__(self, context: Context) -> None:
+        super().__init__(context)
+        self.context.layers.remove(self)
+        self.events = []
+        self.layer = None
+
+    def _handle_event(self, event: mevents.Event):
+        self.events.append(event)
+
+        if not isinstance(event, mevents.DataReceived):
+            # We only ask if we have received new data.
+            return
+
+        yield commands.Hook("next_layer", self)
+
+        # Has an addon decided on the next layer yet?
+        if self.layer:
+            for e in self.events:
+                yield from self.layer.handle_event(e)
+            self.events.clear()
+
+            self._handle_event = self.layer.handle_event
+
+    # Utility methods for addon.
+
+    def data_client(self):
+        return self._data(self.context.client)
+
+    def data_server(self):
+        return self._data(self.context.server)
+
+    def _data(self, connection: Connection):
+        data = (
+            e.data for e in self.events
+            if isinstance(e, mevents.DataReceived) and e.connection == connection
+        )
+        return b"".join(data)

mitmproxy/proxy2/layers/__init__.py

@@ -0,0 +1,5 @@
+from . import modes
+from .http import HTTPLayer
+from .tcp import TCPLayer
+from .tls import TLSLayer
+from .websocket import WebsocketLayer

mitmproxy/proxy2/layers/modes.py

@@ -1,17 +1,13 @@
-from mitmproxy.proxy2 import events
-from mitmproxy.proxy2.commands import TCommandGenerator
+from mitmproxy.net import server_spec
+from mitmproxy.proxy2 import layer
 from mitmproxy.proxy2.context import Context, Server
-from mitmproxy.proxy2.layer import Layer
-from mitmproxy.proxy2.layers.http import HTTPLayer
 
 
-class ReverseProxy(Layer):
-    def __init__(self, context: Context, server_addr):
+class ReverseProxy(layer.Layer):
+    def __init__(self, context: Context):
         super().__init__(context)
+        server_addr = server_spec.parse_with_mode(context.options.mode)[1].address
         self.context.server = Server(server_addr)
-        # self.child_layer = TLSLayer(self.context, True, True)
-        # self.child_layer = TCPLayer(self.context, False)
-        self.child_layer = HTTPLayer(self.context)
 
-    def _handle_event(self, event: events.Event) -> TCommandGenerator:
-        yield from self.child_layer.handle_event(event)
+        child_layer = layer.NextLayer(self.context)
+        self._handle_event = child_layer.handle_event

mitmproxy/proxy2/layers/tcp.py

@@ -9,7 +9,7 @@ class TCPLayer(Layer):
     """
     Simple TCP layer that just relays messages right now.
     """
-    context: Context = None
+    context: Context
     ignore: bool
     flow: tcp.TCPFlow
 

mitmproxy/proxy2/layers/tls.py

@@ -10,25 +10,23 @@ from OpenSSL import SSL
 
 from mitmproxy.certs import CertStore
 from mitmproxy.proxy.protocol.tls import DEFAULT_CLIENT_CIPHERS
-from mitmproxy.proxy2 import events, commands
+from mitmproxy.proxy2 import events, commands, layer
 from mitmproxy.proxy2.context import Context, Connection
-from mitmproxy.proxy2.layer import Layer
-from mitmproxy.proxy2.layers.tcp import TCPLayer
 from mitmproxy.proxy2.utils import expect
 
 
-class TLSLayer(Layer):
-    context: Context = None
-    client_tls: bool = None  # FIXME: not yet used.
-    server_tls: bool = None
-    child_layer: Layer = None
+class TLSLayer(layer.Layer):
+    client_tls: bool  # FIXME: not yet used.
+    server_tls: bool
+    child_layer: layer.Layer = None
+    tls: MutableMapping[Connection, SSL.Connection]
 
     def __init__(self, context: Context, client_tls: bool, server_tls: bool):
         super().__init__(context)
         self.state = self.start
         self.client_tls = client_tls
         self.server_tls = server_tls
-        self.tls: MutableMapping[Connection, SSL.Connection] = {}
+        self.tls = {}
 
     def _handle_event(self, event: events.Event) -> commands.TCommandGenerator:
         yield from self.state(event)
@@ -97,7 +95,7 @@ class TLSLayer(Layer):
 
             if both_handshakes_done:
                 print("both handshakes done")
-                self.child_layer = TCPLayer(self.context)
+                self.child_layer = layer.NextLayer(self.context)
                 yield from self.child_layer.handle_event(events.Start())
                 self.state = self.relay_messages
                 yield from self.state(events.DataReceived(self.context.server, b""))

mitmproxy/proxy2/server.py

@@ -31,8 +31,10 @@ class ConnectionHandler(metaclass=abc.ABCMeta):
         self.client = Client(addr)
         self.context = Context(self.client, None, options)
 
-        # self.layer = ReverseProxy(self.context, ("localhost", 443))
-        self.layer = ReverseProxy(self.context, ("localhost", 8000))
+        if options.mode.startswith("reverse:"):
+            self.layer = ReverseProxy(self.context)
+        else:
+            raise NotImplementedError("Mode not implemented.")
 
         self.transports = {
             self.client: StreamIO(reader, writer)