add more websocket indicators to mitmproxy and mitmdump

4 files changed, 32 insertions, 20 deletions

diff --git a/mitmproxy/addons/dumper.py b/mitmproxy/addons/dumper.py
index 68d59b2d..29f60cfe 100644
--- a/mitmproxy/addons/dumper.py
+++ b/mitmproxy/addons/dumper.py
@@ -238,6 +238,14 @@ class Dumper:
             if self.flow_detail >= 3:
                 self._echo_message(message)
 
+    def websocket_end(self, f):
+        if self.match(f):
+            self.echo("WebSocket connection closed by {}: {} {}, {}".format(
+                f.close_sender,
+                f.close_code,
+                f.close_message,
+                f.close_reason))
+
     def tcp_error(self, f):
         self.echo(
             "Error in TCP connection to {}: {}".format(
diff --git a/mitmproxy/proxy/protocol/websocket.py b/mitmproxy/proxy/protocol/websocket.py
index 69eb898f..15d9a288 100644
--- a/mitmproxy/proxy/protocol/websocket.py
+++ b/mitmproxy/proxy/protocol/websocket.py
@@ -45,13 +45,6 @@ class WebSocketLayer(base.Layer):
         self.server_frame_buffer = []
 
     def _handle_frame(self, frame, source_conn, other_conn, is_server):
-        # sender = "server" if is_server else "client"
-        # self.log(
-        #     "WebSocket frame received from {}".format(sender),
-        #     "debug",
-        #     [repr(frame)]
-        # )
-
         if frame.header.opcode & 0x8 == 0:
             return self._handle_data_frame(frame, source_conn, other_conn, is_server)
         elif frame.header.opcode in (websockets.OPCODE.PING, websockets.OPCODE.PONG):
@@ -102,20 +95,16 @@ class WebSocketLayer(base.Layer):
         return True
 
     def _handle_close(self, frame, source_conn, other_conn, is_server):
-        code = '(status code missing)'
-        msg = None
-        reason = '(message missing)'
+        self.flow.close_sender = "server" if is_server else "client"
         if len(frame.payload) >= 2:
             code, = struct.unpack('!H', frame.payload[:2])
-            msg = websockets.CLOSE_REASON.get_name(code, default='unknown status code')
+            self.flow.close_code = code
+            self.flow.close_message = websockets.CLOSE_REASON.get_name(code, default='unknown status code')
         if len(frame.payload) > 2:
-            reason = frame.payload[2:]
+            self.flow.close_reason = frame.payload[2:]
 
         other_conn.send(bytes(frame))
 
-        sender = "server" if is_server else "client"
-        self.log("WebSocket connection closed by {}: {} {}, {}".format(sender, code, msg, reason), "info")
-
         # close the connection
         return False
 
@@ -130,6 +119,8 @@ class WebSocketLayer(base.Layer):
 
     def __call__(self):
         self.flow = WebSocketFlow(self.client_conn, self.server_conn, self.handshake_flow, self)
+        self.flow.metadata['websocket_handshake'] = self.handshake_flow
+        self.handshake_flow.metadata['websocket_flow'] = self.flow
         self.channel.ask("websocket_start", self.flow)
 
         client = self.client_conn.connection
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index 99c61825..184038ef 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -450,20 +450,29 @@ class ConsoleMaster(master.Master):
     def websocket_message(self, f):
         super().websocket_message(f)
         message = f.messages[-1]
-        self.add_log(message.info, "info")
-        self.add_log(strutils.bytes_to_escaped_str(message.content), "debug")
+        signals.add_log(message.info, "info")
+        signals.add_log(strutils.bytes_to_escaped_str(message.content), "debug")
+
+    @controller.handler
+    def websocket_end(self, f):
+        super().websocket_end(f)
+        signals.add_log("WebSocket connection closed by {}: {} {}, {}".format(
+            f.close_sender,
+            f.close_code,
+            f.close_message,
+            f.close_reason), "info")
 
     @controller.handler
     def tcp_message(self, f):
         super().tcp_message(f)
         message = f.messages[-1]
         direction = "->" if message.from_client else "<-"
-        self.add_log("{client} {direction} tcp {direction} {server}".format(
+        signals.add_log("{client} {direction} tcp {direction} {server}".format(
             client=repr(f.client_conn.address),
             server=repr(f.server_conn.address),
             direction=direction,
         ), "info")
-        self.add_log(strutils.bytes_to_escaped_str(message.content), "debug")
+        signals.add_log(strutils.bytes_to_escaped_str(message.content), "debug")
 
     @controller.handler
     def log(self, evt):
diff --git a/mitmproxy/websocket.py b/mitmproxy/websocket.py
index ef51a392..6e998a52 100644
--- a/mitmproxy/websocket.py
+++ b/mitmproxy/websocket.py
@@ -65,6 +65,10 @@ class WebSocketFlow(flow.Flow):
     def __init__(self, client_conn, server_conn, handshake_flow, live=None):
         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_message = '(message missing)'
+        self.close_reason = 'unknown status code'
         self.handshake_flow = handshake_flow
         self.client_key = websockets.get_client_key(self.handshake_flow.request.headers)
         self.client_protocol = websockets.get_protocol(self.handshake_flow.request.headers)
@@ -80,4 +84,4 @@ class WebSocketFlow(flow.Flow):
     )
 
     def __repr__(self):
-        return "<WebSocketFlow ({} messages)>".format(len(self.messages))
+        return "WebSocketFlow ({} messages)".format(len(self.messages))