diff --git a/mitmproxy/flowfilter.py b/mitmproxy/flowfilter.py index 3a44efd13..304670dbe 100644 --- a/mitmproxy/flowfilter.py +++ b/mitmproxy/flowfilter.py @@ -37,6 +37,7 @@ import sys import functools from mitmproxy import http +from mitmproxy import websocket from mitmproxy import tcp from mitmproxy import flow @@ -99,6 +100,14 @@ class FHTTP(_Action): return True +class FWebSocket(_Action): + code = "websocket" + help = "Match WebSocket flows" + + @only(websocket.WebSocketFlow) + def __call__(self, f): + return True + class FTCP(_Action): code = "tcp" help = "Match TCP flows" @@ -245,7 +254,7 @@ class FBod(_Rex): help = "Body" flags = re.DOTALL - @only(http.HTTPFlow, tcp.TCPFlow) + @only(http.HTTPFlow, websocket.WebSocketFlow, tcp.TCPFlow) def __call__(self, f): if isinstance(f, http.HTTPFlow): if f.request and f.request.raw_content: @@ -254,7 +263,7 @@ class FBod(_Rex): if f.response and f.response.raw_content: if self.re.search(f.response.get_content(strict=False)): return True - elif isinstance(f, tcp.TCPFlow): + elif isinstance(f, websocket.WebSocketFlow) or isinstance(f, tcp.TCPFlow): for msg in f.messages: if self.re.search(msg.content): return True @@ -266,13 +275,13 @@ class FBodRequest(_Rex): help = "Request body" flags = re.DOTALL - @only(http.HTTPFlow, tcp.TCPFlow) + @only(http.HTTPFlow, websocket.WebSocketFlow, tcp.TCPFlow) def __call__(self, f): if isinstance(f, http.HTTPFlow): if f.request and f.request.raw_content: if self.re.search(f.request.get_content(strict=False)): return True - elif isinstance(f, tcp.TCPFlow): + elif isinstance(f, websocket.WebSocketFlow) or isinstance(f, tcp.TCPFlow): for msg in f.messages: if msg.from_client and self.re.search(msg.content): return True @@ -283,13 +292,13 @@ class FBodResponse(_Rex): help = "Response body" flags = re.DOTALL - @only(http.HTTPFlow, tcp.TCPFlow) + @only(http.HTTPFlow, websocket.WebSocketFlow, tcp.TCPFlow) def __call__(self, f): if isinstance(f, http.HTTPFlow): if f.response and f.response.raw_content: if self.re.search(f.response.get_content(strict=False)): return True - elif isinstance(f, tcp.TCPFlow): + elif isinstance(f, websocket.WebSocketFlow) or isinstance(f, tcp.TCPFlow): for msg in f.messages: if not msg.from_client and self.re.search(msg.content): return True diff --git a/mitmproxy/proxy/protocol/websocket.py b/mitmproxy/proxy/protocol/websocket.py index 15d9a288d..d0b12540a 100644 --- a/mitmproxy/proxy/protocol/websocket.py +++ b/mitmproxy/proxy/protocol/websocket.py @@ -59,30 +59,56 @@ class WebSocketLayer(base.Layer): fb.append(frame) if frame.header.fin: - if frame.header.opcode == websockets.OPCODE.TEXT: + payload = b''.join(f.payload for f in fb) + original_chunk_sizes = [len(f.payload) for f in fb] + message_type = fb[0].header.opcode + compressed_message = fb[0].header.rsv1 + fb.clear() + + if message_type == websockets.OPCODE.TEXT: t = WebSocketTextMessage else: t = WebSocketBinaryMessage - payload = b''.join(f.payload for f in fb) - fb.clear() - websocket_message = t(self.flow, not is_server, payload) + length = len(websocket_message.content) self.flow.messages.append(websocket_message) self.channel.ask("websocket_message", self.flow) - # chunk payload into multiple 10kB frames, and send them - payload = websocket_message.content - chunk_size = 10240 # 10kB - chunks = range(0, len(payload), chunk_size) + def get_chunk(payload): + if len(payload) == length: + # message has the same length, we can reuse the same sizes + pos = 0 + for s in original_chunk_sizes: + yield payload[pos:pos + s] + pos += s + else: + # just re-chunk everything into 10kB frames + chunk_size = 10240 + chunks = range(0, len(payload), chunk_size) + for i in chunks: + yield payload[i:i + chunk_size] + frms = [ websockets.Frame( - payload=payload[i:i + chunk_size], + payload=chunk, opcode=frame.header.opcode, mask=(False if is_server else 1), - masking_key=(b'' if is_server else os.urandom(4))) for i in chunks + masking_key=(b'' if is_server else os.urandom(4))) + for chunk in get_chunk(websocket_message.content) ] - frms[-1].header.fin = 1 + + if len(frms) > 0: + frms[-1].header.fin = True + else: + frms.append(websockets.Frame( + fin=True, + opcode=websockets.OPCODE.CONTINUE, + mask=(False if is_server else 1), + masking_key=(b'' if is_server else os.urandom(4)))) + + frms[0].header.opcode = message_type + frms[0].header.rsv1 = compressed_message for frm in frms: other_conn.send(bytes(frm)) @@ -105,7 +131,7 @@ class WebSocketLayer(base.Layer): other_conn.send(bytes(frame)) - # close the connection + # initiate close handshake return False def _handle_unknown_frame(self, frame, source_conn, other_conn, is_server): @@ -126,10 +152,11 @@ class WebSocketLayer(base.Layer): client = self.client_conn.connection server = self.server_conn.connection conns = [client, server] + close_received = False try: while not self.channel.should_exit.is_set(): - r = tcp.ssl_read_select(conns, 1) + r = tcp.ssl_read_select(conns, 0.1) for conn in r: source_conn = self.client_conn if conn == client else self.server_conn other_conn = self.server_conn if conn == client else self.client_conn @@ -137,10 +164,15 @@ class WebSocketLayer(base.Layer): frame = websockets.Frame.from_file(source_conn.rfile) - if not self._handle_frame(frame, source_conn, other_conn, is_server): - return + cont = self._handle_frame(frame, source_conn, other_conn, is_server) + if not cont: + if close_received: + return + else: + close_received = True except (socket.error, exceptions.TcpException, SSL.Error) as e: - self.flow.error = flow.Error("WebSocket connection closed unexpectedly: {}".format(repr(e))) + s = 'server' if is_server else 'client' + self.flow.error = flow.Error("WebSocket connection closed unexpectedly by {}: {}".format(s, repr(e))) self.channel.tell("websocket_error", self.flow) finally: self.channel.tell("websocket_end", self.flow) diff --git a/test/mitmproxy/protocol/test_websocket.py b/test/mitmproxy/protocol/test_websocket.py index e1c3e49a6..e42250e01 100644 --- a/test/mitmproxy/protocol/test_websocket.py +++ b/test/mitmproxy/protocol/test_websocket.py @@ -276,6 +276,7 @@ class TestClose(_WebSocketTest): def handle_websockets(cls, rfile, wfile): frame = websockets.Frame.from_file(rfile) wfile.write(bytes(frame)) + wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE))) wfile.flush() with pytest.raises(exceptions.TcpDisconnect): @@ -287,6 +288,7 @@ class TestClose(_WebSocketTest): client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE))) client.wfile.flush() + websockets.Frame.from_file(client.rfile) with pytest.raises(exceptions.TcpDisconnect): websockets.Frame.from_file(client.rfile) @@ -296,6 +298,7 @@ class TestClose(_WebSocketTest): client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE, payload=b'\00\42'))) client.wfile.flush() + websockets.Frame.from_file(client.rfile) with pytest.raises(exceptions.TcpDisconnect): websockets.Frame.from_file(client.rfile) @@ -305,6 +308,7 @@ class TestClose(_WebSocketTest): client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE, payload=b'\00\42foobar'))) client.wfile.flush() + websockets.Frame.from_file(client.rfile) with pytest.raises(exceptions.TcpDisconnect): websockets.Frame.from_file(client.rfile)