1 from __future__
import annotations
11 create_socks_proxy_socket
,
12 make_socks_proxy_opts
,
15 from .common
import Features
, Response
, register_rh
16 from .exceptions
import (
17 CertificateVerifyError
,
24 from .websocket
import WebSocketRequestHandler
, WebSocketResponse
25 from ..compat
import functools
26 from ..dependencies
import websockets
27 from ..socks
import ProxyError
as SocksProxyError
28 from ..utils
import int_or_none
31 raise ImportError('websockets is not installed')
33 import websockets
.version
35 websockets_version
= tuple(map(int_or_none
, websockets
.version
.version
.split('.')))
36 if websockets_version
< (12, 0):
37 raise ImportError('Only websockets>=12.0 is supported')
39 import websockets
.sync
.client
40 from websockets
.uri
import parse_uri
42 # In websockets Connection, recv_exc and recv_events_exc are defined
43 # after the recv events handler thread is started [1].
44 # On our CI using PyPy, in some cases a race condition may occur
45 # where the recv events handler thread tries to use these attributes before they are defined [2].
46 # 1: https://github.com/python-websockets/websockets/blame/de768cf65e7e2b1a3b67854fb9e08816a5ff7050/src/websockets/sync/connection.py#L93
47 # 2: "AttributeError: 'ClientConnection' object has no attribute 'recv_events_exc'. Did you mean: 'recv_events'?"
48 import websockets
.sync
.connection
# isort: split
49 with contextlib
.suppress(Exception):
51 websockets
.sync
.connection
.Connection
.recv_exc
= None
53 websockets
.sync
.connection
.Connection
.recv_events_exc
= None
56 class WebsocketsResponseAdapter(WebSocketResponse
):
58 def __init__(self
, ws
: websockets
.sync
.client
.ClientConnection
, url
):
60 fp
=io
.BytesIO(ws
.response
.body
or b
''),
62 headers
=ws
.response
.headers
,
63 status
=ws
.response
.status_code
,
64 reason
=ws
.response
.reason_phrase
,
72 def send(self
, message
):
73 # https://websockets.readthedocs.io/en/stable/reference/sync/client.html#websockets.sync.client.ClientConnection.send
75 return self
._ws
.send(message
)
76 except (websockets
.exceptions
.WebSocketException
, RuntimeError, TimeoutError
) as e
:
77 raise TransportError(cause
=e
) from e
78 except SocksProxyError
as e
:
79 raise ProxyError(cause
=e
) from e
80 except TypeError as e
:
81 raise RequestError(cause
=e
) from e
84 # https://websockets.readthedocs.io/en/stable/reference/sync/client.html#websockets.sync.client.ClientConnection.recv
86 return self
._ws
.recv()
87 except SocksProxyError
as e
:
88 raise ProxyError(cause
=e
) from e
89 except (websockets
.exceptions
.WebSocketException
, RuntimeError, TimeoutError
) as e
:
90 raise TransportError(cause
=e
) from e
94 class WebsocketsRH(WebSocketRequestHandler
):
96 Websockets request handler
97 https://websockets.readthedocs.io
98 https://github.com/python-websockets/websockets
100 _SUPPORTED_URL_SCHEMES
= ('wss', 'ws')
101 _SUPPORTED_PROXY_SCHEMES
= ('socks4', 'socks4a', 'socks5', 'socks5h')
102 _SUPPORTED_FEATURES
= (Features
.ALL_PROXY
, Features
.NO_PROXY
)
103 RH_NAME
= 'websockets'
105 def __init__(self
, *args
, **kwargs
):
106 super().__init
__(*args
, **kwargs
)
107 self
.__logging
_handlers
= {}
108 for name
in ('websockets.client', 'websockets.server'):
109 logger
= logging
.getLogger(name
)
110 handler
= logging
.StreamHandler(stream
=sys
.stdout
)
111 handler
.setFormatter(logging
.Formatter(f
'{self.RH_NAME}: %(message)s'))
112 self
.__logging
_handlers
[name
] = handler
113 logger
.addHandler(handler
)
115 logger
.setLevel(logging
.DEBUG
)
117 def _check_extensions(self
, extensions
):
118 super()._check
_extensions
(extensions
)
119 extensions
.pop('timeout', None)
120 extensions
.pop('cookiejar', None)
123 # Remove the logging handler that contains a reference to our logger
124 # See: https://github.com/yt-dlp/yt-dlp/issues/8922
125 for name
, handler
in self
.__logging
_handlers
.items():
126 logging
.getLogger(name
).removeHandler(handler
)
128 def _send(self
, request
):
129 timeout
= self
._calculate
_timeout
(request
)
130 headers
= self
._merge
_headers
(request
.headers
)
131 if 'cookie' not in headers
:
132 cookiejar
= self
._get
_cookiejar
(request
)
133 cookie_header
= cookiejar
.get_cookie_header(request
.url
)
135 headers
['cookie'] = cookie_header
137 wsuri
= parse_uri(request
.url
)
138 create_conn_kwargs
= {
139 'source_address': (self
.source_address
, 0) if self
.source_address
else None,
142 proxy
= select_proxy(request
.url
, self
._get
_proxies
(request
))
145 socks_proxy_options
= make_socks_proxy_opts(proxy
)
146 sock
= create_connection(
147 address
=(socks_proxy_options
['addr'], socks_proxy_options
['port']),
148 _create_socket_func
=functools
.partial(
149 create_socks_proxy_socket
, (wsuri
.host
, wsuri
.port
), socks_proxy_options
),
153 sock
= create_connection(
154 address
=(wsuri
.host
, wsuri
.port
),
157 conn
= websockets
.sync
.client
.connect(
160 additional_headers
=headers
,
161 open_timeout
=timeout
,
162 user_agent_header
=None,
163 ssl_context
=self
._make
_sslcontext
() if wsuri
.secure
else None,
164 close_timeout
=0, # not ideal, but prevents yt-dlp hanging
166 return WebsocketsResponseAdapter(conn
, url
=request
.url
)
168 # Exceptions as per https://websockets.readthedocs.io/en/stable/reference/sync/client.html
169 except SocksProxyError
as e
:
170 raise ProxyError(cause
=e
) from e
171 except websockets
.exceptions
.InvalidURI
as e
:
172 raise RequestError(cause
=e
) from e
173 except ssl
.SSLCertVerificationError
as e
:
174 raise CertificateVerifyError(cause
=e
) from e
175 except ssl
.SSLError
as e
:
176 raise SSLError(cause
=e
) from e
177 except websockets
.exceptions
.InvalidStatus
as e
:
180 fp
=io
.BytesIO(e
.response
.body
),
182 headers
=e
.response
.headers
,
183 status
=e
.response
.status_code
,
184 reason
=e
.response
.reason_phrase
),
186 except (OSError, TimeoutError
, websockets
.exceptions
.WebSocketException
) as e
:
187 raise TransportError(cause
=e
) from e