[sans-io] merge ClientServerContext and Context

mitmproxy/proxy2/context.py

@@ -1,3 +1,6 @@
+from typing import Optional
+
+
 class Connection:
     """
     Connections exposed to the layers only contain metadata, no socket objects.
@@ -33,18 +36,12 @@ class Context:
     """
 
     client: Client
+    server: Optional[Server]
 
-    def __init__(self, client: Client) -> None:
+    def __init__(
+            self,
+            client: Client,
+            server: Optional[Server],
+    ) -> None:
         self.client = client
-
-
-class ClientServerContext(Context):
-    """
-    In most cases, there's also only exactly one server.
-    """
-
-    server: Server
-
-    def __init__(self, client: Client, server: Server) -> None:
-        super().__init__(client)
         self.server = server

mitmproxy/proxy2/layers/http.py

@@ -7,7 +7,7 @@ import h11
 from mitmproxy.net import http
 from mitmproxy.net import websockets
 from mitmproxy.proxy2 import events, commands
-from mitmproxy.proxy2.context import ClientServerContext
+from mitmproxy.proxy2.context import Context
 from mitmproxy.proxy2.layer import Layer
 from mitmproxy.proxy2.layers import websocket
 from mitmproxy.proxy2.utils import expect
@@ -17,12 +17,12 @@ class HTTPLayer(Layer):
     """
     Simple TCP layer that just relays messages right now.
     """
-    context: ClientServerContext = None
+    context: Context = None
 
     # this is like a mini state machine.
     state: typing.Callable[[events.Event], commands.TCommandGenerator]
 
-    def __init__(self, context: ClientServerContext):
+    def __init__(self, context: Context):
         super().__init__(context)
         self.state = self.read_request_headers
         self.flow = mock.Mock()

mitmproxy/proxy2/layers/modes.py

@@ -1,6 +1,6 @@
 from mitmproxy.proxy2 import events
 from mitmproxy.proxy2.commands import TCommandGenerator
-from mitmproxy.proxy2.context import Context, Server, ClientServerContext
+from mitmproxy.proxy2.context import Context, Server
 from mitmproxy.proxy2.layer import Layer
 from mitmproxy.proxy2.layers.http import HTTPLayer
 
@@ -8,11 +8,10 @@ from mitmproxy.proxy2.layers.http import HTTPLayer
 class ReverseProxy(Layer):
     def __init__(self, context: Context, server_addr):
         super().__init__(context)
-        server = Server(server_addr)
-        self.child_context = ClientServerContext(context.client, server)
-        # self.child_layer = TLSLayer(self.child_context, True, True)
-        # self.child_layer = TCPLayer(self.child_context, False)
-        self.child_layer = HTTPLayer(self.child_context)
+        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)

mitmproxy/proxy2/layers/tcp.py

@@ -1,6 +1,6 @@
 from mitmproxy import tcp, flow
 from mitmproxy.proxy2 import commands, events
-from mitmproxy.proxy2.context import ClientServerContext
+from mitmproxy.proxy2.context import Context
 from mitmproxy.proxy2.layer import Layer
 from mitmproxy.proxy2.utils import expect
 
@@ -9,11 +9,11 @@ class TCPLayer(Layer):
     """
     Simple TCP layer that just relays messages right now.
     """
-    context: ClientServerContext = None
+    context: Context = None
     ignore: bool
     flow: tcp.TCPFlow
 
-    def __init__(self, context: ClientServerContext, ignore: bool = False):
+    def __init__(self, context: Context, ignore: bool = False):
         super().__init__(context)
         self.ignore = ignore
         self.flow = None

mitmproxy/proxy2/layers/tls.py

@@ -11,19 +11,19 @@ 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.context import ClientServerContext, Connection
+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: ClientServerContext = None
+    context: Context = None
     client_tls: bool = None  # FIXME: not yet used.
     server_tls: bool = None
     child_layer: Layer = None
 
-    def __init__(self, context: ClientServerContext, client_tls: bool, server_tls: bool):
+    def __init__(self, context: Context, client_tls: bool, server_tls: bool):
         super().__init__(context)
         self.state = self.start
         self.client_tls = client_tls

mitmproxy/proxy2/layers/websocket.py

@@ -1,6 +1,6 @@
 from mitmproxy import websocket, http, flow
 from mitmproxy.proxy2 import events, commands
-from mitmproxy.proxy2.context import ClientServerContext
+from mitmproxy.proxy2.context import Context
 from mitmproxy.proxy2.layer import Layer
 from mitmproxy.proxy2.utils import expect
 from wsproto import connection as wsconn
@@ -11,10 +11,10 @@ class WebsocketLayer(Layer):
     """
     Simple TCP layer that just relays messages right now.
     """
-    context: ClientServerContext = None
+    context: Context = None
     flow: websocket.WebSocketFlow
 
-    def __init__(self, context: ClientServerContext, handshake_flow: http.HTTPFlow):
+    def __init__(self, context: Context, handshake_flow: http.HTTPFlow):
         super().__init__(context)
         self.flow = websocket.WebSocketFlow(context.client, context.server, handshake_flow)
         assert context.server.connected

mitmproxy/proxy2/server.py

@@ -28,7 +28,7 @@ class ConnectionHandler(metaclass=abc.ABCMeta):
         addr = writer.get_extra_info('peername')
 
         self.client = Client(addr)
-        self.context = Context(self.client)
+        self.context = Context(self.client, None)
 
         # self.layer = ReverseProxy(self.context, ("localhost", 443))
         self.layer = ReverseProxy(self.context, ("localhost", 8000))

mitmproxy/proxy2/test/conftest.py

@@ -5,7 +5,7 @@ from mitmproxy.proxy2 import context
 
 @pytest.fixture
 def tctx():
-    return context.ClientServerContext(
+    return context.Context(
         context.Client("client"),
         context.Server("server")
     )