From 7e04ac0578b60821711249993659f3a1a7b48584 Mon Sep 17 00:00:00 2001
From: Ujjwal Verma <ujjwalverma1111@gmail.com>
Date: Tue, 4 Jul 2017 10:50:34 +0200
Subject: add websocket streaming option

---
 mitmproxy/addons/streambodies.py           | 15 ++++++---------
 mitmproxy/options.py                       |  7 +++++++
 test/mitmproxy/addons/test_streambodies.py |  1 +
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/mitmproxy/addons/streambodies.py b/mitmproxy/addons/streambodies.py
index b98ed1fa..c841075f 100644
--- a/mitmproxy/addons/streambodies.py
+++ b/mitmproxy/addons/streambodies.py
@@ -2,7 +2,6 @@ from mitmproxy.net.http import http1
 from mitmproxy import exceptions
 from mitmproxy import ctx
 from mitmproxy.utils import human
-from mitmproxy import websocket
 
 
 class StreamBodies:
@@ -18,13 +17,6 @@ class StreamBodies:
 
     def run(self, f, is_request):
         if self.max_size:
-            if isinstance(f, websocket.WebSocketFlow):
-                f.stream = True
-                ctx.log.info("Streaming WebSocket message {client} - {server}".format(
-                    client=human.format_address(f.client_conn.address),
-                    server=human.format_address(f.server_conn.address))
-                )
-                return
             r = f.request if is_request else f.response
             try:
                 expected_size = http1.expected_http_body_size(
@@ -45,4 +37,9 @@ class StreamBodies:
         self.run(f, False)
 
     def websocket_start(self, f):
-        self.run(f, False)
+        if ctx.options.stream_websockets:
+            f.stream = True
+            ctx.log.info("Streaming WebSocket messages between {client} and {server}".format(
+                client=human.format_address(f.client_conn.address),
+                server=human.format_address(f.server_conn.address))
+            )
diff --git a/mitmproxy/options.py b/mitmproxy/options.py
index a3872679..e6c2fed6 100644
--- a/mitmproxy/options.py
+++ b/mitmproxy/options.py
@@ -154,6 +154,13 @@ class Options(optmanager.OptManager):
             Understands k/m/g suffixes, i.e. 3m for 3 megabytes.
             """
         )
+        self.add_option(
+            "stream_websockets", bool, False,
+            """
+            Stream WebSocket messages between client and server.
+            Messages are captured and cannot be modified.
+            """
+        )
         self.add_option(
             "verbosity", int, 2,
             "Log verbosity."
diff --git a/test/mitmproxy/addons/test_streambodies.py b/test/mitmproxy/addons/test_streambodies.py
index 54799949..426ec9ae 100644
--- a/test/mitmproxy/addons/test_streambodies.py
+++ b/test/mitmproxy/addons/test_streambodies.py
@@ -30,6 +30,7 @@ def test_simple():
         f.response.headers["content-length"] = "invalid"
         tctx.cycle(sa, f)
 
+        tctx.configure(sa, stream_websockets = True)
         f = tflow.twebsocketflow()
         assert not f.stream
         sa.websocket_start(f)
-- 
cgit v1.2.3