From ab1bcfe7d2d3724f08d77339ce118033ce37c0a6 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Thu, 11 Mar 2021 10:58:00 +0100 Subject: [PATCH] websocket: incorporate PR feedback --- mitmproxy/addons/dumper.py | 2 +- mitmproxy/io/compat.py | 6 +----- mitmproxy/proxy/layers/http/__init__.py | 8 +++----- mitmproxy/proxy/layers/websocket.py | 2 +- mitmproxy/tools/console/common.py | 15 +++++++++------ mitmproxy/websocket.py | 4 ++-- 6 files changed, 17 insertions(+), 20 deletions(-) diff --git a/mitmproxy/addons/dumper.py b/mitmproxy/addons/dumper.py index 80ae72730..7ebb128b5 100644 --- a/mitmproxy/addons/dumper.py +++ b/mitmproxy/addons/dumper.py @@ -300,7 +300,7 @@ class Dumper: def websocket_end(self, f: http.HTTPFlow): assert f.websocket is not None # satisfy type checker if self.match(f): - c = 'client' if f.websocket.close_by_client else 'server' + c = 'client' if f.websocket.closed_by_client else 'server' self.echo(f"WebSocket connection closed by {c}: {f.websocket.close_code} {f.websocket.close_reason}") def tcp_error(self, f): diff --git a/mitmproxy/io/compat.py b/mitmproxy/io/compat.py index 963048614..16887dd1a 100644 --- a/mitmproxy/io/compat.py +++ b/mitmproxy/io/compat.py @@ -282,17 +282,13 @@ def convert_11_12(data): 'type': 'http', 'version': 12 } - data["request"]["scheme"] = { - b"http": b"ws", - b"https": b"wss" - }.get(data["request"]["scheme"], data["request"]["scheme"]) data["metadata"]["duplicated"] = ( "This WebSocket flow has been migrated from an old file format version " "and may appear duplicated." ) data["websocket"] = { "messages": ws_flow["messages"], - "close_by_client": ws_flow["close_sender"] == "client", + "closed_by_client": ws_flow["close_sender"] == "client", "close_code": ws_flow["close_code"], "close_reason": ws_flow["close_reason"], } diff --git a/mitmproxy/proxy/layers/http/__init__.py b/mitmproxy/proxy/layers/http/__init__.py index c333acd1f..3ba891ea5 100644 --- a/mitmproxy/proxy/layers/http/__init__.py +++ b/mitmproxy/proxy/layers/http/__init__.py @@ -4,6 +4,8 @@ import time from dataclasses import dataclass from typing import DefaultDict, Dict, List, Optional, Tuple, Union +import wsproto.handshake + from mitmproxy import flow, http from mitmproxy.connection import Connection, Server from mitmproxy.net import server_spec @@ -315,7 +317,7 @@ class HttpStream(layer.Layer): and self.flow.response.headers.get("upgrade", "").lower() == "websocket" and - self.flow.request.headers.get("Sec-WebSocket-Version", "") == "13" + self.flow.request.headers.get("Sec-WebSocket-Version", "").encode() == wsproto.handshake.WEBSOCKET_VERSION and self.context.options.websocket ) @@ -323,10 +325,6 @@ class HttpStream(layer.Layer): # We need to set this before calling the response hook # so that addons can determine if a WebSocket connection is following up. self.flow.websocket = WebSocketData() - if self.flow.request.scheme == "http": - self.flow.request.scheme = "ws" - elif self.flow.request.scheme == "https": - self.flow.request.scheme = "wss" yield HttpResponseHook(self.flow) self.server_state = self.state_done diff --git a/mitmproxy/proxy/layers/websocket.py b/mitmproxy/proxy/layers/websocket.py index f39a93ddc..f83ef14f0 100644 --- a/mitmproxy/proxy/layers/websocket.py +++ b/mitmproxy/proxy/layers/websocket.py @@ -171,7 +171,7 @@ class WebsocketLayer(layer.Layer): ) yield dst_ws.send2(ws_event) elif isinstance(ws_event, wsproto.events.CloseConnection): - self.flow.websocket.close_by_client = from_client + self.flow.websocket.closed_by_client = from_client self.flow.websocket.close_code = ws_event.code self.flow.websocket.close_reason = ws_event.reason diff --git a/mitmproxy/tools/console/common.py b/mitmproxy/tools/console/common.py index 910e1af88..edb4a0d1e 100644 --- a/mitmproxy/tools/console/common.py +++ b/mitmproxy/tools/console/common.py @@ -299,12 +299,8 @@ def colorize_url(url): parts = url.split('/', 3) if len(parts) < 4 or len(parts[1]) > 0 or parts[0][-1:] != ':': return [('error', len(url))] # bad URL - schemes = { - 'http:': 'scheme_http', - 'https:': 'scheme_https', - } return [ - (schemes.get(parts[0], "scheme_other"), len(parts[0]) - 1), + (SCHEME_STYLES.get(parts[0], "scheme_other"), len(parts[0]) - 1), ('url_punctuation', 3), # :// ] + colorize_host(parts[2]) + colorize_req('/' + parts[3]) @@ -701,6 +697,13 @@ def format_flow( response_content_type = None duration = None + scheme = f.request.scheme + if f.websocket is not None: + if scheme == "https": + scheme = "wss" + elif scheme == "http": + scheme = "ws" + if render_mode in (RenderMode.LIST, RenderMode.DETAILVIEW): render_func = format_http_flow_list else: @@ -711,7 +714,7 @@ def format_flow( marked=f.marked, is_replay=f.is_replay, request_method=f.request.method, - request_scheme=f.request.scheme, + request_scheme=scheme, request_host=f.request.pretty_host if hostheader else f.request.host, request_path=f.request.path, request_url=f.request.pretty_url if hostheader else f.request.url, diff --git a/mitmproxy/websocket.py b/mitmproxy/websocket.py index f1e37542b..841944058 100644 --- a/mitmproxy/websocket.py +++ b/mitmproxy/websocket.py @@ -95,7 +95,7 @@ class WebSocketData(stateobject.StateObject): messages: List[WebSocketMessage] """All `WebSocketMessage`s transferred over this connection.""" - close_by_client: Optional[bool] = None + closed_by_client: Optional[bool] = None """ True if the client closed the connection, False if the server closed the connection, @@ -108,7 +108,7 @@ class WebSocketData(stateobject.StateObject): _stateobject_attributes = dict( messages=List[WebSocketMessage], - close_by_client=bool, + closed_by_client=bool, close_code=int, close_reason=str, )