From 1a7ce384dac5099308b68e629c88e7b81ad44866 Mon Sep 17 00:00:00 2001
From: Thomas Kriechbaumer <thomas@kriechbaumer.name>
Date: Sun, 17 Dec 2017 18:13:35 +0100
Subject: websocket: support domain and url filters

---
 mitmproxy/flowfilter.py           | 11 +++++++----
 mitmproxy/test/tflow.py           |  2 +-
 test/mitmproxy/test_flowfilter.py | 14 ++++++++++++++
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/mitmproxy/flowfilter.py b/mitmproxy/flowfilter.py
index 23e47e2b..12f2bac8 100644
--- a/mitmproxy/flowfilter.py
+++ b/mitmproxy/flowfilter.py
@@ -322,14 +322,15 @@ class FDomain(_Rex):
     flags = re.IGNORECASE
     is_binary = False
 
-    @only(http.HTTPFlow)
+    @only(http.HTTPFlow, websocket.WebSocketFlow)
     def __call__(self, f):
+        if isinstance(f, websocket.WebSocketFlow):
+            f = f.handshake_flow
         return bool(
             self.re.search(f.request.host) or
             self.re.search(f.request.pretty_host)
         )
 
-
 class FUrl(_Rex):
     code = "u"
     help = "URL"
@@ -342,9 +343,11 @@ class FUrl(_Rex):
             toks = toks[1:]
         return klass(*toks)
 
-    @only(http.HTTPFlow)
+    @only(http.HTTPFlow, websocket.WebSocketFlow)
     def __call__(self, f):
-        if not f.request:
+        if isinstance(f, websocket.WebSocketFlow):
+            f = f.handshake_flow
+        if not f or not f.request:
             return False
         return self.re.search(f.request.pretty_url)
 
diff --git a/mitmproxy/test/tflow.py b/mitmproxy/test/tflow.py
index c3dab30c..cb2b940a 100644
--- a/mitmproxy/test/tflow.py
+++ b/mitmproxy/test/tflow.py
@@ -44,7 +44,7 @@ def twebsocketflow(client_conn=True, server_conn=True, messages=True, err=None,
             "GET",
             "http",
             "example.com",
-            "80",
+            80,
             "/ws",
             "HTTP/1.1",
             headers=net_http.Headers(
diff --git a/test/mitmproxy/test_flowfilter.py b/test/mitmproxy/test_flowfilter.py
index c411258a..4eb37d81 100644
--- a/test/mitmproxy/test_flowfilter.py
+++ b/test/mitmproxy/test_flowfilter.py
@@ -420,6 +420,20 @@ class TestMatchingWebSocketFlow:
         e = self.err()
         assert self.q("~e", e)
 
+    def test_domain(self):
+        q = self.flow()
+        assert self.q("~d example.com", q)
+        assert not self.q("~d none", q)
+
+    def test_url(self):
+        q = self.flow()
+        assert self.q("~u example.com", q)
+        assert self.q("~u example.com/ws", q)
+        assert not self.q("~u moo/path", q)
+
+        q.handshake_flow = None
+        assert not self.q("~u example.com", q)
+
     def test_body(self):
         f = self.flow()
 
-- 
cgit v1.2.3


From 8d836d251e1ce0b7141155fdc84aed8efc1850a6 Mon Sep 17 00:00:00 2001
From: Thomas Kriechbaumer <thomas@kriechbaumer.name>
Date: Sun, 17 Dec 2017 17:44:36 +0000
Subject: fix #2640

---
 mitmproxy/addons/save.py              |  9 +++++++++
 mitmproxy/contrib/wsproto/__init__.py |  0
 mitmproxy/flowfilter.py               |  1 +
 mitmproxy/master.py                   | 25 ++++++++++++++++++++-----
 mitmproxy/proxy/protocol/http.py      |  1 +
 mitmproxy/proxy/protocol/websocket.py | 10 +++++++---
 mitmproxy/websocket.py                | 18 +++++++++++++-----
 test/mitmproxy/addons/test_save.py    | 12 ++++++++++++
 8 files changed, 63 insertions(+), 13 deletions(-)
 create mode 100644 mitmproxy/contrib/wsproto/__init__.py

diff --git a/mitmproxy/addons/save.py b/mitmproxy/addons/save.py
index 40cd6f82..7bb55a84 100644
--- a/mitmproxy/addons/save.py
+++ b/mitmproxy/addons/save.py
@@ -74,6 +74,15 @@ class Save:
             self.stream.add(flow)
             self.active_flows.discard(flow)
 
+    def websocket_start(self, flow):
+        if self.stream:
+            self.active_flows.add(flow)
+
+    def websocket_end(self, flow):
+        if self.stream:
+            self.stream.add(flow)
+            self.active_flows.discard(flow)
+
     def response(self, flow):
         if self.stream:
             self.stream.add(flow)
diff --git a/mitmproxy/contrib/wsproto/__init__.py b/mitmproxy/contrib/wsproto/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/mitmproxy/flowfilter.py b/mitmproxy/flowfilter.py
index 12f2bac8..d1fd8299 100644
--- a/mitmproxy/flowfilter.py
+++ b/mitmproxy/flowfilter.py
@@ -331,6 +331,7 @@ class FDomain(_Rex):
             self.re.search(f.request.pretty_host)
         )
 
+
 class FUrl(_Rex):
     code = "u"
     help = "URL"
diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index 5997ff6d..9660c137 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -9,6 +9,7 @@ from mitmproxy import eventsequence
 from mitmproxy import exceptions
 from mitmproxy import command
 from mitmproxy import http
+from mitmproxy import websocket
 from mitmproxy import log
 from mitmproxy.net import server_spec
 from mitmproxy.proxy.protocol import http_replay
@@ -41,6 +42,7 @@ class Master:
         self.should_exit = threading.Event()
         self._server = None
         self.first_tick = True
+        self.waiting_flows = []
 
     @property
     def server(self):
@@ -117,15 +119,28 @@ class Master:
         self.should_exit.set()
         self.addons.trigger("done")
 
+    def _change_reverse_host(self, f):
+        if self.options.mode.startswith("reverse:"):
+            _, upstream_spec = server_spec.parse_with_mode(self.options.mode)
+            f.request.host, f.request.port = upstream_spec.address
+            f.request.scheme = upstream_spec.scheme
+
     def load_flow(self, f):
         """
-        Loads a flow
+        Loads a flow and links websocket & handshake flows
         """
+
         if isinstance(f, http.HTTPFlow):
-            if self.options.mode.startswith("reverse:"):
-                _, upstream_spec = server_spec.parse_with_mode(self.options.mode)
-                f.request.host, f.request.port = upstream_spec.address
-                f.request.scheme = upstream_spec.scheme
+            self._change_reverse_host(f)
+            if 'websocket' in f.metadata:
+                self.waiting_flows.append(f)
+
+        if isinstance(f, websocket.WebSocketFlow):
+            hf = [hf for hf in self.waiting_flows if hf.id == f.metadata['websocket_handshake']][0]
+            f.handshake_flow = hf
+            self.waiting_flows.remove(hf)
+            self._change_reverse_host(f.handshake_flow)
+
         f.reply = controller.DummyReply()
         for e, o in eventsequence.iterate(f):
             self.addons.handle_lifecycle(e, o)
diff --git a/mitmproxy/proxy/protocol/http.py b/mitmproxy/proxy/protocol/http.py
index 57ac0f16..076ffa62 100644
--- a/mitmproxy/proxy/protocol/http.py
+++ b/mitmproxy/proxy/protocol/http.py
@@ -321,6 +321,7 @@ class HttpLayer(base.Layer):
 
         try:
             if websockets.check_handshake(request.headers) and websockets.check_client_version(request.headers):
+                f.metadata['websocket'] = True
                 # We only support RFC6455 with WebSocket version 13
                 # allow inline scripts to manipulate the client handshake
                 self.channel.ask("websocket_handshake", f)
diff --git a/mitmproxy/proxy/protocol/websocket.py b/mitmproxy/proxy/protocol/websocket.py
index 34dcba06..1bd5284d 100644
--- a/mitmproxy/proxy/protocol/websocket.py
+++ b/mitmproxy/proxy/protocol/websocket.py
@@ -1,10 +1,11 @@
 import socket
 from OpenSSL import SSL
 
+
+from mitmproxy.contrib import wsproto
 from mitmproxy.contrib.wsproto import events
 from mitmproxy.contrib.wsproto.connection import ConnectionType, WSConnection
 from mitmproxy.contrib.wsproto.extensions import PerMessageDeflate
-from mitmproxy.contrib.wsproto.frame_protocol import Opcode
 
 from mitmproxy import exceptions
 from mitmproxy import flow
@@ -93,11 +94,14 @@ class WebSocketLayer(base.Layer):
 
         if event.message_finished:
             original_chunk_sizes = [len(f) for f in fb]
-            message_type = Opcode.TEXT if isinstance(event, events.TextReceived) else Opcode.BINARY
-            if message_type == Opcode.TEXT:
+
+            if isinstance(event, events.TextReceived):
+                message_type = wsproto.frame_protocol.Opcode.TEXT
                 payload = ''.join(fb)
             else:
+                message_type = wsproto.frame_protocol.Opcode.BINARY
                 payload = b''.join(fb)
+
             fb.clear()
 
             websocket_message = WebSocketMessage(message_type, not is_server, payload)
diff --git a/mitmproxy/websocket.py b/mitmproxy/websocket.py
index 6c1e7000..ade23732 100644
--- a/mitmproxy/websocket.py
+++ b/mitmproxy/websocket.py
@@ -1,6 +1,8 @@
 import time
 from typing import List, Optional
 
+from mitmproxy.contrib import wsproto
+
 from mitmproxy import flow
 from mitmproxy.net import websockets
 from mitmproxy.coretypes import serializable
@@ -11,7 +13,7 @@ class WebSocketMessage(serializable.Serializable):
     def __init__(
         self, type: int, from_client: bool, content: bytes, timestamp: Optional[int]=None
     ) -> None:
-        self.type = type
+        self.type = wsproto.frame_protocol.Opcode(type)
         self.from_client = from_client
         self.content = content
         self.timestamp = timestamp or int(time.time())  # type: int
@@ -21,13 +23,14 @@ class WebSocketMessage(serializable.Serializable):
         return cls(*state)
 
     def get_state(self):
-        return self.type, self.from_client, self.content, self.timestamp
+        return int(self.type), self.from_client, self.content, self.timestamp
 
     def set_state(self, state):
         self.type, self.from_client, self.content, self.timestamp = state
+        self.type = wsproto.frame_protocol.Opcode(self.type)  # replace enum with bare int
 
     def __repr__(self):
-        if self.type == websockets.OPCODE.TEXT:
+        if self.type == wsproto.frame_protocol.Opcode.TEXT:
             return "text message: {}".format(repr(self.content))
         else:
             return "binary message: {}".format(strutils.bytes_to_escaped_str(self.content))
@@ -42,7 +45,7 @@ class WebSocketFlow(flow.Flow):
         super().__init__("websocket", client_conn, server_conn, live)
         self.messages = []  # type: List[WebSocketMessage]
         self.close_sender = 'client'
-        self.close_code = '(status code missing)'
+        self.close_code = wsproto.frame_protocol.CloseReason.NORMAL_CLOSURE
         self.close_message = '(message missing)'
         self.close_reason = 'unknown status code'
         self.stream = False
@@ -69,7 +72,7 @@ class WebSocketFlow(flow.Flow):
     _stateobject_attributes.update(dict(
         messages=List[WebSocketMessage],
         close_sender=str,
-        close_code=str,
+        close_code=int,
         close_message=str,
         close_reason=str,
         client_key=str,
@@ -83,6 +86,11 @@ class WebSocketFlow(flow.Flow):
         # dumping the handshake_flow will include the WebSocketFlow too.
     ))
 
+    def get_state(self):
+        d = super().get_state()
+        d['close_code'] = int(d['close_code'])  # replace enum with bare int
+        return d
+
     @classmethod
     def from_state(cls, state):
         f = cls(None, None, None)
diff --git a/test/mitmproxy/addons/test_save.py b/test/mitmproxy/addons/test_save.py
index a4e425cd..84564157 100644
--- a/test/mitmproxy/addons/test_save.py
+++ b/test/mitmproxy/addons/test_save.py
@@ -44,6 +44,18 @@ def test_tcp(tmpdir):
         assert rd(p)
 
 
+def test_websocket(tmpdir):
+    sa = save.Save()
+    with taddons.context() as tctx:
+        p = str(tmpdir.join("foo"))
+        tctx.configure(sa, save_stream_file=p)
+
+        f = tflow.twebsocketflow()
+        sa.websocket_start(f)
+        tctx.configure(sa, save_stream_file=None)
+        assert rd(p)
+
+
 def test_save_command(tmpdir):
     sa = save.Save()
     with taddons.context() as tctx:
-- 
cgit v1.2.3


From 4a3702792cc0270773d321b7b28daa0e3df023c2 Mon Sep 17 00:00:00 2001
From: Thomas Kriechbaumer <thomas@kriechbaumer.name>
Date: Mon, 18 Dec 2017 09:33:14 +0100
Subject: ignore contrib type errors

---
 mitmproxy/contrib/wsproto/extensions.py     | 2 ++
 mitmproxy/contrib/wsproto/frame_protocol.py | 2 ++
 setup.cfg                                   | 3 +++
 3 files changed, 7 insertions(+)

diff --git a/mitmproxy/contrib/wsproto/extensions.py b/mitmproxy/contrib/wsproto/extensions.py
index f7cf4fb6..0e0d2018 100644
--- a/mitmproxy/contrib/wsproto/extensions.py
+++ b/mitmproxy/contrib/wsproto/extensions.py
@@ -1,3 +1,5 @@
+# type: ignore
+
 # -*- coding: utf-8 -*-
 """
 wsproto/extensions
diff --git a/mitmproxy/contrib/wsproto/frame_protocol.py b/mitmproxy/contrib/wsproto/frame_protocol.py
index b95dceec..30f146c6 100644
--- a/mitmproxy/contrib/wsproto/frame_protocol.py
+++ b/mitmproxy/contrib/wsproto/frame_protocol.py
@@ -1,3 +1,5 @@
+# type: ignore
+
 # -*- coding: utf-8 -*-
 """
 wsproto/frame_protocol
diff --git a/setup.cfg b/setup.cfg
index fd31d15b..7c754722 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -19,6 +19,9 @@ exclude_lines =
     pragma: no cover
     raise NotImplementedError()
 
+[mypy-mitmproxy.contrib.*]
+ignore_errors = True
+
 [tool:full_coverage]
 exclude =
     mitmproxy/proxy/protocol/base.py
-- 
cgit v1.2.3


From 417a860f322769961dfae470aa065df4457162e1 Mon Sep 17 00:00:00 2001
From: Thomas Kriechbaumer <thomas@kriechbaumer.name>
Date: Mon, 18 Dec 2017 10:17:14 +0100
Subject: add tests

---
 mitmproxy/contrib/wsproto/__init__.py | 13 +++++++++++++
 mitmproxy/test/tflow.py               |  4 +++-
 mitmproxy/websocket.py                |  2 +-
 test/mitmproxy/addons/test_save.py    |  1 +
 test/mitmproxy/test_flow.py           | 16 +++++++++++++++-
 5 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/mitmproxy/contrib/wsproto/__init__.py b/mitmproxy/contrib/wsproto/__init__.py
index e69de29b..d0592bc5 100644
--- a/mitmproxy/contrib/wsproto/__init__.py
+++ b/mitmproxy/contrib/wsproto/__init__.py
@@ -0,0 +1,13 @@
+from . import compat
+from . import connection
+from . import events
+from . import extensions
+from . import frame_protocol
+
+__all__ = [
+    'compat',
+    'connection',
+    'events',
+    'extensions',
+    'frame_protocol',
+]
diff --git a/mitmproxy/test/tflow.py b/mitmproxy/test/tflow.py
index cb2b940a..91747866 100644
--- a/mitmproxy/test/tflow.py
+++ b/mitmproxy/test/tflow.py
@@ -75,7 +75,9 @@ def twebsocketflow(client_conn=True, server_conn=True, messages=True, err=None,
         handshake_flow.response = resp
 
     f = websocket.WebSocketFlow(client_conn, server_conn, handshake_flow)
-    handshake_flow.metadata['websocket_flow'] = f
+    f.metadata['websocket_handshake'] = handshake_flow.id
+    handshake_flow.metadata['websocket_flow'] = f.id
+    handshake_flow.metadata['websocket'] = True
 
     if messages is True:
         messages = [
diff --git a/mitmproxy/websocket.py b/mitmproxy/websocket.py
index ade23732..8efd4117 100644
--- a/mitmproxy/websocket.py
+++ b/mitmproxy/websocket.py
@@ -13,7 +13,7 @@ class WebSocketMessage(serializable.Serializable):
     def __init__(
         self, type: int, from_client: bool, content: bytes, timestamp: Optional[int]=None
     ) -> None:
-        self.type = wsproto.frame_protocol.Opcode(type)
+        self.type = wsproto.frame_protocol.Opcode(type)  # type: ignore
         self.from_client = from_client
         self.content = content
         self.timestamp = timestamp or int(time.time())  # type: int
diff --git a/test/mitmproxy/addons/test_save.py b/test/mitmproxy/addons/test_save.py
index 84564157..2dee708f 100644
--- a/test/mitmproxy/addons/test_save.py
+++ b/test/mitmproxy/addons/test_save.py
@@ -52,6 +52,7 @@ def test_websocket(tmpdir):
 
         f = tflow.twebsocketflow()
         sa.websocket_start(f)
+        sa.websocket_end(f)
         tctx.configure(sa, save_stream_file=None)
         assert rd(p)
 
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index fcc766b5..8cc11a16 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -97,7 +97,7 @@ class TestSerialize:
 
 
 class TestFlowMaster:
-    def test_load_flow_reverse(self):
+    def test_load_http_flow_reverse(self):
         s = tservers.TestState()
         opts = options.Options(
             mode="reverse:https://use-this-domain"
@@ -108,6 +108,20 @@ class TestFlowMaster:
         fm.load_flow(f)
         assert s.flows[0].request.host == "use-this-domain"
 
+    def test_load_websocket_flow(self):
+        s = tservers.TestState()
+        opts = options.Options(
+            mode="reverse:https://use-this-domain"
+        )
+        fm = master.Master(opts)
+        fm.addons.add(s)
+        f = tflow.twebsocketflow()
+        fm.load_flow(f.handshake_flow)
+        fm.load_flow(f)
+        assert s.flows[0].request.host == "use-this-domain"
+        assert s.flows[1].handshake_flow == f.handshake_flow
+        assert len(s.flows[1].messages) == len(f.messages)
+
     def test_replay(self):
         opts = options.Options()
         fm = master.Master(opts)
-- 
cgit v1.2.3


From 05a8559dff7a04d8c4109f98357321031c0981f9 Mon Sep 17 00:00:00 2001
From: Thomas Kriechbaumer <thomas@kriechbaumer.name>
Date: Mon, 18 Dec 2017 19:51:29 +0100
Subject: add doc string

---
 mitmproxy/master.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index 9660c137..de3b24e1 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -120,6 +120,11 @@ class Master:
         self.addons.trigger("done")
 
     def _change_reverse_host(self, f):
+        """
+        When we load flows in reverse proxy mode, we adjust the target host to
+        the reverse proxy destination for all flows we load. This makes it very
+        easy to replay saved flows against a different host.
+        """
         if self.options.mode.startswith("reverse:"):
             _, upstream_spec = server_spec.parse_with_mode(self.options.mode)
             f.request.host, f.request.port = upstream_spec.address
-- 
cgit v1.2.3