Merge pull request #1566 from mhils/issue-1546

add websocket on/off switch, improve logging (fix #1546, fix #1547)

7 files changed, 66 insertions, 46 deletions

diff --git a/mitmproxy/cmdline.py b/mitmproxy/cmdline.py
index 9fb4a561..ff431909 100644
--- a/mitmproxy/cmdline.py
+++ b/mitmproxy/cmdline.py
@@ -1,14 +1,12 @@
 from __future__ import absolute_import, print_function, division
 
+import configargparse
 import os
 import re
-
-import configargparse
-
 from mitmproxy import exceptions
 from mitmproxy import filt
-from mitmproxy import platform
 from mitmproxy import options
+from mitmproxy import platform
 from netlib import human
 from netlib import tcp
 from netlib import version
@@ -257,6 +255,7 @@ def get_common_options(args):
         no_upstream_cert = args.no_upstream_cert,
         spoof_source_address = args.spoof_source_address,
         rawtcp = args.rawtcp,
+        websockets = args.websockets,
         upstream_server = upstream_server,
         upstream_auth = args.upstream_auth,
         ssl_version_client = args.ssl_version_client,
@@ -475,6 +474,13 @@ def proxy_options(parser):
                         "Disabled by default. "
                         "Default value will change in a future version."
                         )
+    websockets = group.add_mutually_exclusive_group()
+    websockets.add_argument("--websockets", action="store_true", dest="websockets")
+    websockets.add_argument("--no-websockets", action="store_false", dest="websockets",
+                            help="Explicitly enable/disable experimental WebSocket support. "
+                                 "Disabled by default as messages are only printed to the event log and not retained. "
+                                 "Default value will change in a future version."
+                            )
     group.add_argument(
         "--spoof-source-address",
         action="store_true", dest="spoof_source_address",
diff --git a/mitmproxy/options.py b/mitmproxy/options.py
index ba4ed0c7..ccf06f75 100644
--- a/mitmproxy/options.py
+++ b/mitmproxy/options.py
@@ -70,6 +70,7 @@ class Options(optmanager.OptManager):
             mode = "regular",  # type: str
             no_upstream_cert = False,  # type: bool
             rawtcp = False,  # type: bool
+            websockets = False,  # type: bool
             spoof_source_address = False,  # type: bool
             upstream_server = "",  # type: str
             upstream_auth = "",  # type: str
@@ -129,6 +130,7 @@ class Options(optmanager.OptManager):
         self.mode = mode
         self.no_upstream_cert = no_upstream_cert
         self.rawtcp = rawtcp
+        self.websockets = websockets
         self.spoof_source_address = spoof_source_address
         self.upstream_server = upstream_server
         self.upstream_auth = upstream_auth
diff --git a/mitmproxy/protocol/http.py b/mitmproxy/protocol/http.py
index 1632e66f..894ae465 100644
--- a/mitmproxy/protocol/http.py
+++ b/mitmproxy/protocol/http.py
@@ -1,17 +1,15 @@
 from __future__ import absolute_import, print_function, division
 
-import time
-import sys
-import traceback
-
 import h2.exceptions
+import netlib.exceptions
 import six
-
+import sys
+import time
+import traceback
 from mitmproxy import exceptions
 from mitmproxy import models
 from mitmproxy.protocol import base
-
-import netlib.exceptions
+from mitmproxy.protocol import websockets as pwebsockets
 from netlib import http
 from netlib import tcp
 from netlib import websockets
@@ -193,8 +191,8 @@ class HttpLayer(base.Layer):
 
             try:
                 if websockets.check_handshake(request.headers) and websockets.check_client_version(request.headers):
-                    # we only support RFC6455 with WebSockets version 13
-                    # allow inline scripts to manupulate the client handshake
+                    # We only support RFC6455 with WebSockets version 13
+                    # allow inline scripts to manipulate the client handshake
                     self.channel.ask("websockets_handshake", flow)
 
                 if not flow.response:
@@ -217,12 +215,8 @@ class HttpLayer(base.Layer):
                     return
 
                 # Handle 101 Switching Protocols
-                # It may be useful to pass additional args (such as the upgrade header)
-                # to next_layer in the future
                 if flow.response.status_code == 101:
-                    layer = self.ctx.next_layer(self, flow)
-                    layer()
-                    return
+                    return self.handle_101_switching_protocols(flow)
 
                 # Upstream Proxy Mode: Handle CONNECT
                 if flow.request.first_line_format == "authority" and flow.response.status_code == 200:
@@ -438,3 +432,26 @@ class HttpLayer(base.Layer):
                     ))
                 return False
         return True
+
+    def handle_101_switching_protocols(self, flow):
+        """
+        Handle a successful HTTP 101 Switching Protocols Response, received after e.g. a WebSocket upgrade request.
+        """
+        # Check for WebSockets handshake
+        is_websockets = (
+            flow and
+            websockets.check_handshake(flow.request.headers) and
+            websockets.check_handshake(flow.response.headers)
+        )
+        if is_websockets and not self.config.options.websockets:
+            self.log(
+                "Client requested WebSocket connection, but the protocol is currently disabled in mitmproxy.",
+                "info"
+            )
+
+        if is_websockets and self.config.options.websockets:
+            layer = pwebsockets.WebSocketsLayer(self, flow)
+        else:
+            layer = self.ctx.next_layer(self)
+
+        layer()
diff --git a/mitmproxy/protocol/websockets.py b/mitmproxy/protocol/websockets.py
index f15a38ef..7ac386f1 100644
--- a/mitmproxy/protocol/websockets.py
+++ b/mitmproxy/protocol/websockets.py
@@ -1,14 +1,12 @@
 from __future__ import absolute_import, print_function, division
 
+import netlib.exceptions
 import socket
 import struct
-
 from OpenSSL import SSL
-
 from mitmproxy import exceptions
 from mitmproxy.protocol import base
-
-import netlib.exceptions
+from netlib import strutils
 from netlib import tcp
 from netlib import websockets
 
@@ -49,22 +47,28 @@ class WebSocketsLayer(base.Layer):
         self.server_extensions = websockets.get_extensions(self._flow.response.headers)
 
     def _handle_frame(self, frame, source_conn, other_conn, is_server):
+        sender = "server" if is_server else "client"
         self.log(
-            "WebSockets Frame received from {}".format("server" if is_server else "client"),
+            "WebSockets Frame received from {}".format(sender),
             "debug",
             [repr(frame)]
         )
 
         if frame.header.opcode & 0x8 == 0:
+            self.log(
+                "{direction} websocket {direction} {server}".format(
+                    server=repr(self.server_conn.address),
+                    direction="<-" if is_server else "->",
+                ),
+                "info",
+                strutils.bytes_to_escaped_str(frame.payload, keep_spacing=True).splitlines()
+            )
             # forward the data frame to the other side
             other_conn.send(bytes(frame))
-            self.log("WebSockets frame received by {}: {}".format(is_server, frame), "debug")
         elif frame.header.opcode in (websockets.OPCODE.PING, websockets.OPCODE.PONG):
             # just forward the ping/pong to the other side
             other_conn.send(bytes(frame))
         elif frame.header.opcode == websockets.OPCODE.CLOSE:
-            other_conn.send(bytes(frame))
-
             code = '(status code missing)'
             msg = None
             reason = '(message missing)'
@@ -73,11 +77,13 @@ class WebSocketsLayer(base.Layer):
                 msg = websockets.CLOSE_REASON.get_name(code, default='unknown status code')
             if len(frame.payload) > 2:
                 reason = frame.payload[2:]
-            self.log("WebSockets connection closed: {} {}, {}".format(code, msg, reason), "info")
+            self.log("WebSockets connection closed by {}: {} {}, {}".format(sender, code, msg, reason), "info")
 
+            other_conn.send(bytes(frame))
             # close the connection
             return False
         else:
+            self.log("Unknown WebSockets frame received from {}".format(sender), "info", [repr(frame)])
             # unknown frame - just forward it
             other_conn.send(bytes(frame))
 
diff --git a/mitmproxy/proxy/config.py b/mitmproxy/proxy/config.py
index 2cf8410a..df4dd4c8 100644
--- a/mitmproxy/proxy/config.py
+++ b/mitmproxy/proxy/config.py
@@ -10,6 +10,7 @@ import six
 from OpenSSL import SSL, crypto
 
 from mitmproxy import exceptions
+from mitmproxy import options as moptions  # noqa
 from netlib import certutils
 from netlib import tcp
 from netlib.http import authentication
@@ -71,7 +72,7 @@ def parse_upstream_auth(auth):
 class ProxyConfig:
 
     def __init__(self, options):
-        self.options = options
+        self.options = options  # type: moptions.Options
 
         self.authenticator = None
         self.check_ignore = None
@@ -83,7 +84,7 @@ class ProxyConfig:
         options.changed.connect(self.configure)
 
     def configure(self, options, updated):
-        # type: (mitmproxy.options.Options, Any) -> None
+        # type: (moptions.Options, Any) -> None
         if options.add_upstream_certs_to_client_chain and not options.ssl_insecure:
             raise exceptions.OptionsError(
                 "The verify-upstream-cert requires certificate verification to be disabled. "
diff --git a/mitmproxy/proxy/root_context.py b/mitmproxy/proxy/root_context.py
index 95611362..81dd625c 100644
--- a/mitmproxy/proxy/root_context.py
+++ b/mitmproxy/proxy/root_context.py
@@ -4,7 +4,6 @@ import sys
 
 import six
 
-from netlib import websockets
 import netlib.exceptions
 from mitmproxy import exceptions
 from mitmproxy import protocol
@@ -33,7 +32,7 @@ class RootContext(object):
         self.channel = channel
         self.config = config
 
-    def next_layer(self, top_layer, flow=None):
+    def next_layer(self, top_layer):
         """
         This function determines the next layer in the protocol stack.
 
@@ -43,22 +42,10 @@ class RootContext(object):
         Returns:
             The next layer
         """
-        layer = self._next_layer(top_layer, flow)
+        layer = self._next_layer(top_layer)
         return self.channel.ask("next_layer", layer)
 
-    def _next_layer(self, top_layer, flow):
-        if flow is not None:
-            # We already have a flow, try to derive the next information from it
-
-            # Check for WebSockets handshake
-            is_websockets = (
-                flow and
-                websockets.check_handshake(flow.request.headers) and
-                websockets.check_handshake(flow.response.headers)
-            )
-            if isinstance(top_layer, protocol.HttpLayer) and is_websockets:
-                return protocol.WebSocketsLayer(top_layer, flow)
-
+    def _next_layer(self, top_layer):
         try:
             d = top_layer.client_conn.rfile.peek(3)
         except netlib.exceptions.TcpException as e:
diff --git a/test/mitmproxy/protocol/test_websockets.py b/test/mitmproxy/protocol/test_websockets.py
index e7e2684f..e2361d89 100644
--- a/test/mitmproxy/protocol/test_websockets.py
+++ b/test/mitmproxy/protocol/test_websockets.py
@@ -65,7 +65,8 @@ class _WebSocketsTestBase(object):
         opts = options.Options(
             listen_port=0,
             no_upstream_cert=False,
-            ssl_insecure=True
+            ssl_insecure=True,
+            websockets=True,
         )
         opts.cadir = os.path.join(tempfile.gettempdir(), "mitmproxy")
         return opts