From 976cec10c953853d27f8c72f5dd0f2a195a27bf7 Mon Sep 17 00:00:00 2001
From: coletdjnz <coletdjnz@protonmail.com>
Date: Sun, 19 May 2024 10:17:28 +1200
Subject: [PATCH] Fix debug traffic for websockets proxy

---
 yt_dlp/networking/_helper.py     |  3 +++
 yt_dlp/networking/_websockets.py | 21 ++++++++++++++++++++-
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/yt_dlp/networking/_helper.py b/yt_dlp/networking/_helper.py
index 4bfe29c06c..1a448997eb 100644
--- a/yt_dlp/networking/_helper.py
+++ b/yt_dlp/networking/_helper.py
@@ -307,6 +307,7 @@ def create_http_connect_connection(
     source_address=None,
     username=None,
     password=None,
+    debug=False,
 ):
 
     proxy_headers = dict()
@@ -316,6 +317,8 @@ def create_http_connect_connection(
             f'{username or ""}:{password or ""}'.encode('utf-8')).decode('utf-8')
 
     conn = HTTPConnection(proxy_host, port=proxy_port, timeout=timeout)
+    conn.set_debuglevel(int(debug))
+
     conn.response_class = NoCloseHTTPResponse
 
     if hasattr(conn, '_create_connection'):
diff --git a/yt_dlp/networking/_websockets.py b/yt_dlp/networking/_websockets.py
index d10ab33844..8874811efa 100644
--- a/yt_dlp/networking/_websockets.py
+++ b/yt_dlp/networking/_websockets.py
@@ -4,7 +4,6 @@
 import io
 import logging
 import ssl
-import sys
 import urllib.parse
 
 from ._helper import (
@@ -70,6 +69,25 @@
     websockets.sync.connection.Connection.recv_events_exc = None
 
 
+class WebsocketsLoggingHandler(logging.Handler):
+    """Redirect websocket logs to our logger"""
+
+    def __init__(self, logger, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._logger = logger
+
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            if record.levelno >= logging.ERROR:
+                self._logger.error(msg)
+            else:
+                self._logger.stdout(msg)
+
+        except Exception:
+            self.handleError(record)
+
+
 class WebsocketsResponseAdapter(WebSocketResponse):
 
     def __init__(self, ws: websockets.sync.client.ClientConnection, url):
@@ -181,6 +199,7 @@ def _make_sock(self, proxy, url, timeout):
                     source_address=self.source_address,
                     username=parsed_proxy_url.username,
                     password=parsed_proxy_url.password,
+                    debug=self.verbose,
                 )
         return create_connection(
             address=(parsed_url.host, parsed_url.port),