websocket: incorporate PR feedback

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):

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"],
         }

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

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
 

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,

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,
     )