diff options
Diffstat (limited to 'www/mitmproxy/files/patch-mitmproxy_proxy_protocol_websocket.py')
| -rw-r--r-- | www/mitmproxy/files/patch-mitmproxy_proxy_protocol_websocket.py | 174 |
1 files changed, 174 insertions, 0 deletions
diff --git a/www/mitmproxy/files/patch-mitmproxy_proxy_protocol_websocket.py b/www/mitmproxy/files/patch-mitmproxy_proxy_protocol_websocket.py new file mode 100644 index 000000000000..b0f07b3f54a5 --- /dev/null +++ b/www/mitmproxy/files/patch-mitmproxy_proxy_protocol_websocket.py @@ -0,0 +1,174 @@ +--- mitmproxy/proxy/protocol/websocket.py.orig 2019-02-14 21:03:58 UTC ++++ mitmproxy/proxy/protocol/websocket.py +@@ -4,8 +4,9 @@ from OpenSSL import SSL + + + import wsproto +-from wsproto import events +-from wsproto.connection import ConnectionType, WSConnection ++from wsproto import events, WSConnection ++from wsproto.connection import ConnectionType ++from wsproto.events import AcceptConnection, CloseConnection, Message, Ping, Request + from wsproto.extensions import PerMessageDeflate + + from mitmproxy import exceptions +@@ -52,51 +53,52 @@ class WebSocketLayer(base.Layer): + + self.connections: dict[object, WSConnection] = {} + +- extensions = [] ++ client_extensions = [] ++ server_extensions = [] + if 'Sec-WebSocket-Extensions' in handshake_flow.response.headers: + if PerMessageDeflate.name in handshake_flow.response.headers['Sec-WebSocket-Extensions']: +- extensions = [PerMessageDeflate()] +- self.connections[self.client_conn] = WSConnection(ConnectionType.SERVER, +- extensions=extensions) +- self.connections[self.server_conn] = WSConnection(ConnectionType.CLIENT, +- host=handshake_flow.request.host, +- resource=handshake_flow.request.path, +- extensions=extensions) +- if extensions: +- for conn in self.connections.values(): +- conn.extensions[0].finalize(conn, handshake_flow.response.headers['Sec-WebSocket-Extensions']) ++ client_extensions = [PerMessageDeflate()] ++ server_extensions = [PerMessageDeflate()] ++ self.connections[self.client_conn] = WSConnection(ConnectionType.SERVER) ++ self.connections[self.server_conn] = WSConnection(ConnectionType.CLIENT) + +- data = self.connections[self.server_conn].bytes_to_send() +- self.connections[self.client_conn].receive_bytes(data) ++ if client_extensions: ++ client_extensions[0].finalize(handshake_flow.response.headers['Sec-WebSocket-Extensions']) ++ if server_extensions: ++ server_extensions[0].finalize(handshake_flow.response.headers['Sec-WebSocket-Extensions']) + ++ request = Request(extensions=client_extensions, host=handshake_flow.request.host, target=handshake_flow.request.path) ++ data = self.connections[self.server_conn].send(request) ++ self.connections[self.client_conn].receive_data(data) ++ + event = next(self.connections[self.client_conn].events()) +- assert isinstance(event, events.ConnectionRequested) ++ assert isinstance(event, events.Request) + +- self.connections[self.client_conn].accept(event) +- self.connections[self.server_conn].receive_bytes(self.connections[self.client_conn].bytes_to_send()) +- assert isinstance(next(self.connections[self.server_conn].events()), events.ConnectionEstablished) ++ data = self.connections[self.client_conn].send(AcceptConnection(extensions=server_extensions)) ++ self.connections[self.server_conn].receive_data(data) ++ assert isinstance(next(self.connections[self.server_conn].events()), events.AcceptConnection) + + def _handle_event(self, event, source_conn, other_conn, is_server): +- if isinstance(event, events.DataReceived): +- return self._handle_data_received(event, source_conn, other_conn, is_server) +- elif isinstance(event, events.PingReceived): +- return self._handle_ping_received(event, source_conn, other_conn, is_server) +- elif isinstance(event, events.PongReceived): +- return self._handle_pong_received(event, source_conn, other_conn, is_server) +- elif isinstance(event, events.ConnectionClosed): +- return self._handle_connection_closed(event, source_conn, other_conn, is_server) ++ if isinstance(event, events.Message): ++ return self._handle_message(event, source_conn, other_conn, is_server) ++ elif isinstance(event, events.Ping): ++ return self._handle_ping(event, source_conn, other_conn, is_server) ++ elif isinstance(event, events.Pong): ++ return self._handle_pong(event, source_conn, other_conn, is_server) ++ elif isinstance(event, events.CloseConnection): ++ return self._handle_close_connection(event, source_conn, other_conn, is_server) + + # fail-safe for unhandled events + return True # pragma: no cover + +- def _handle_data_received(self, event, source_conn, other_conn, is_server): ++ def _handle_message(self, event, source_conn, other_conn, is_server): + fb = self.server_frame_buffer if is_server else self.client_frame_buffer + fb.append(event.data) + + if event.message_finished: + original_chunk_sizes = [len(f) for f in fb] + +- if isinstance(event, events.TextReceived): ++ if isinstance(event, events.TextMessage): + message_type = wsproto.frame_protocol.Opcode.TEXT + payload = ''.join(fb) + else: +@@ -127,19 +129,20 @@ class WebSocketLayer(base.Layer): + yield (payload[i:i + chunk_size], True if i + chunk_size >= len(payload) else False) + + for chunk, final in get_chunk(websocket_message.content): +- self.connections[other_conn].send_data(chunk, final) +- other_conn.send(self.connections[other_conn].bytes_to_send()) ++ data = self.connections[other_conn].send(Message(data=chunk, message_finished=final)) ++ other_conn.send(data) + + if self.flow.stream: +- self.connections[other_conn].send_data(event.data, event.message_finished) +- other_conn.send(self.connections[other_conn].bytes_to_send()) ++ data = self.connections[other_conn].send(Message(data=event.data, message_finished=event.message_finished)) ++ other_conn.send(data) + return True + +- def _handle_ping_received(self, event, source_conn, other_conn, is_server): +- # PING is automatically answered with a PONG by wsproto +- self.connections[other_conn].ping() +- other_conn.send(self.connections[other_conn].bytes_to_send()) +- source_conn.send(self.connections[source_conn].bytes_to_send()) ++ def _handle_ping(self, event, source_conn, other_conn, is_server): ++ # Use event.response to create the approprate Pong response ++ data = self.connections[other_conn].send(Ping()) ++ other_conn.send(data) ++ data = self.connections[source_conn].send(event.response()) ++ source_conn.send(data) + self.log( + "Ping Received from {}".format("server" if is_server else "client"), + "info", +@@ -147,7 +150,7 @@ class WebSocketLayer(base.Layer): + ) + return True + +- def _handle_pong_received(self, event, source_conn, other_conn, is_server): ++ def _handle_pong(self, event, source_conn, other_conn, is_server): + self.log( + "Pong Received from {}".format("server" if is_server else "client"), + "info", +@@ -155,14 +158,15 @@ class WebSocketLayer(base.Layer): + ) + return True + +- def _handle_connection_closed(self, event, source_conn, other_conn, is_server): ++ def _handle_close_connection(self, event, source_conn, other_conn, is_server): + self.flow.close_sender = "server" if is_server else "client" + self.flow.close_code = event.code + self.flow.close_reason = event.reason + +- self.connections[other_conn].close(event.code, event.reason) +- other_conn.send(self.connections[other_conn].bytes_to_send()) +- source_conn.send(self.connections[source_conn].bytes_to_send()) ++ data = self.connections[other_conn].send(CloseConnection(code=event.code, reason=event.reason)) ++ other_conn.send(data) ++ data = self.connections[source_conn].send(event.response()) ++ source_conn.send(data) + + return False + +@@ -170,8 +174,7 @@ class WebSocketLayer(base.Layer): + while True: + try: + payload = message_queue.get_nowait() +- self.connections[endpoint].send_data(payload, final=True) +- data = self.connections[endpoint].bytes_to_send() ++ data = self.connections[endpoint].send(Message(data=payload, message_finished=True)) + endpoint.send(data) + except queue.Empty: + break +@@ -197,8 +200,8 @@ class WebSocketLayer(base.Layer): + is_server = (source_conn == self.server_conn) + + frame = websockets.Frame.from_file(source_conn.rfile) +- self.connections[source_conn].receive_bytes(bytes(frame)) +- source_conn.send(self.connections[source_conn].bytes_to_send()) ++ data = self.connections[source_conn].receive_data(bytes(frame)) ++ source_conn.send(data) + + if close_received: + return |