add websocket stub

mitmproxy/addons/proxyserver.py

@@ -83,5 +83,7 @@ class Proxyserver:
             self.loop.call_soon_threadsafe(lambda: asyncio.ensure_future(self.start()))
 
     def request(self, flow):
-        print("Changing port...")
-        ctx.options.listen_port += 1
+        pass
+        # test live options changes.
+        # print("Changing port...")
+        # ctx.options.listen_port += 1

mitmproxy/proxy/protocol2/commands.py

@@ -8,6 +8,7 @@ The counterpart to commands are events.
 """
 import typing
 
+from mitmproxy import log
 from mitmproxy.proxy.protocol2.context import Connection
 
 
@@ -79,4 +80,9 @@ class Hook(Command):
         self.data = data
 
 
+class Log(Hook):
+    def __init__(self, *args, level="info"):
+        super().__init__("log", log.LogEntry(repr(args), level))
+
+
 TCommandGenerator = typing.Generator[Command, typing.Any, None]

mitmproxy/proxy/protocol2/http.py

@@ -1,8 +1,10 @@
 import typing
 from warnings import warn
 
+import sys
+
 import h11
-from mitmproxy.proxy.protocol2 import events, commands
+from mitmproxy.proxy.protocol2 import events, commands, websocket
 from mitmproxy.proxy.protocol2.context import ClientServerContext
 from mitmproxy.proxy.protocol2.layer import Layer
 from mitmproxy.proxy.protocol2.utils import expect
@@ -28,7 +30,7 @@ class HTTPLayer(Layer):
         def log_event(orig):
             def next_event():
                 e = orig()
-                print(e)
+                print(e, file=sys.__stdout__)
                 return e
 
             return next_event
@@ -56,7 +58,7 @@ class HTTPLayer(Layer):
         if event is h11.NEED_DATA:
             return
         elif isinstance(event, h11.Request):
-            yield commands.Hook("requestheaders", event)
+            yield commands.Log("requestheaders", event)
 
             if self.client_conn.client_is_waiting_for_100_continue:
                 raise NotImplementedError()
@@ -76,7 +78,7 @@ class HTTPLayer(Layer):
                 self.flow_events[0].append(event)
             elif isinstance(event, h11.EndOfMessage):
                 self.flow_events[0].append(event)
-                yield commands.Hook("request", self.flow_events)
+                yield commands.Log("request", self.flow_events)
                 yield from self._send_request()
                 return
             else:
@@ -84,7 +86,12 @@ class HTTPLayer(Layer):
 
     def _send_request(self):
         if not self.context.server.connected:
-            yield commands.OpenConnection(self.context.server)
+            err = yield commands.OpenConnection(self.context.server)
+            if err:
+                yield commands.Log("error", err)
+                yield commands.CloseConnection(self.context.client)
+                self._handle_event = self.done
+                return
         for e in self.flow_events[0]:
             bytes_to_send = self.server_conn.send(e)
             yield commands.SendData(self.context.server, bytes_to_send)
@@ -95,11 +102,18 @@ class HTTPLayer(Layer):
         if event is h11.NEED_DATA:
             return
         elif isinstance(event, h11.Response):
-            yield commands.Hook("responseheaders", event)
+            yield commands.Log("responseheaders", event)
 
             self.flow_events[1].append(event)
             self.state = self.read_response_body
             yield from self.read_response_body()  # there may already be further events.
+        elif isinstance(event, h11.InformationalResponse):
+            if event.status_code == 101:
+                # FIXME: check if this is actually WebSocket
+                child_layer = websocket.WebsocketLayer(self.context, None)
+                yield from child_layer.handle_event(events.Start())
+                self._handle_event = child_layer.handle_event
+                return
         else:
             raise TypeError(f"Unexpected event: {event}")
 
@@ -112,7 +126,7 @@ class HTTPLayer(Layer):
                 self.flow_events[1].append(event)
             elif isinstance(event, h11.EndOfMessage):
                 self.flow_events[1].append(event)
-                yield commands.Hook("response", self.flow_events)
+                yield commands.Log("response", self.flow_events)
                 yield from self._send_response()
                 return
             else:
@@ -128,3 +142,7 @@ class HTTPLayer(Layer):
         self.flow_events = [[], []]
         self.client_conn.start_next_cycle()
         self.server_conn.start_next_cycle()
+
+    @expect(events.DataReceived, events.ConnectionClosed)
+    def done(self, _):
+        yield from ()

mitmproxy/proxy/protocol2/layer.py

@@ -8,6 +8,7 @@ from abc import ABCMeta, abstractmethod
 from mitmproxy.proxy.protocol2 import commands, events
 from mitmproxy.proxy.protocol2.context import Context
 from mitmproxy.proxy.protocol2.events import Event
+from mitmproxy.proxy.protocol2.utils import expect
 
 
 class Paused(typing.NamedTuple):

mitmproxy/proxy/protocol2/reverse_proxy.py

@@ -12,8 +12,8 @@ class ReverseProxy(Layer):
         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.child_layer = TCPLayer(self.child_context, False)
+        self.child_layer = HTTPLayer(self.child_context)
 
     def _handle_event(self, event: Event) -> TCommandGenerator:
         yield from self.child_layer.handle_event(event)

mitmproxy/proxy/protocol2/server/server_async.py

@@ -32,7 +32,7 @@ class ConnectionHandler(metaclass=abc.ABCMeta):
         self.context = Context(self.client)
 
         # self.layer = ReverseProxy(self.context, ("localhost", 443))
-        self.layer = ReverseProxy(self.context, ("localhost", 80))
+        self.layer = ReverseProxy(self.context, ("localhost", 8000))
 
         self.transports = {
             self.client: StreamIO(reader, writer)

mitmproxy/proxy/protocol2/websocket.py

@@ -0,0 +1,29 @@
+from mitmproxy import websocket, http
+from mitmproxy.proxy.protocol2 import events, commands
+from mitmproxy.proxy.protocol2.context import ClientServerContext
+from mitmproxy.proxy.protocol2.layer import Layer
+from mitmproxy.proxy.protocol2.utils import expect
+
+
+class WebsocketLayer(Layer):
+    """
+    Simple TCP layer that just relays messages right now.
+    """
+    context: ClientServerContext = None
+    flow: websocket.WebSocketFlow
+
+    def __init__(self, context: ClientServerContext, handshake_flow: http.HTTPFlow):
+        super().__init__(context)
+        self.flow = websocket.WebSocketFlow(context.client, context.server, handshake_flow)
+        assert context.server.connected
+
+    @expect(events.Start)
+    def start(self, _) -> commands.TCommandGenerator:
+        yield from ()
+        self._handle_event = self.relay_messages
+
+    @expect(events.DataReceived, events.ConnectionClosed)
+    def relay_messages(self, event: events.Event) -> commands.TCommandGenerator:
+        raise NotImplementedError()
+
+    _handle_event = start