1 from __future__
import annotations
10 create_socks_proxy_socket
,
11 make_socks_proxy_opts
,
14 from .common
import Features
, Response
, register_rh
15 from .exceptions
import (
16 CertificateVerifyError
,
23 from .websocket
import WebSocketRequestHandler
, WebSocketResponse
24 from ..compat
import functools
25 from ..dependencies
import websockets
26 from ..socks
import ProxyError
as SocksProxyError
27 from ..utils
import int_or_none
30 raise ImportError('websockets is not installed')
32 import websockets
.version
34 websockets_version
= tuple(map(int_or_none
, websockets
.version
.version
.split('.')))
35 if websockets_version
< (12, 0):
36 raise ImportError('Only websockets>=12.0 is supported')
38 import websockets
.sync
.client
39 from websockets
.uri
import parse_uri
42 class WebsocketsResponseAdapter(WebSocketResponse
):
44 def __init__(self
, wsw
: websockets
.sync
.client
.ClientConnection
, url
):
46 fp
=io
.BytesIO(wsw
.response
.body
or b
''),
48 headers
=wsw
.response
.headers
,
49 status
=wsw
.response
.status_code
,
50 reason
=wsw
.response
.reason_phrase
,
58 def send(self
, message
):
59 # https://websockets.readthedocs.io/en/stable/reference/sync/client.html#websockets.sync.client.ClientConnection.send
61 return self
.wsw
.send(message
)
62 except (websockets
.exceptions
.WebSocketException
, RuntimeError, TimeoutError
) as e
:
63 raise TransportError(cause
=e
) from e
64 except SocksProxyError
as e
:
65 raise ProxyError(cause
=e
) from e
66 except TypeError as e
:
67 raise RequestError(cause
=e
) from e
70 # https://websockets.readthedocs.io/en/stable/reference/sync/client.html#websockets.sync.client.ClientConnection.recv
72 return self
.wsw
.recv()
73 except SocksProxyError
as e
:
74 raise ProxyError(cause
=e
) from e
75 except (websockets
.exceptions
.WebSocketException
, RuntimeError, TimeoutError
) as e
:
76 raise TransportError(cause
=e
) from e
80 class WebsocketsRH(WebSocketRequestHandler
):
82 Websockets request handler
83 https://websockets.readthedocs.io
84 https://github.com/python-websockets/websockets
86 _SUPPORTED_URL_SCHEMES
= ('wss', 'ws')
87 _SUPPORTED_PROXY_SCHEMES
= ('socks4', 'socks4a', 'socks5', 'socks5h')
88 _SUPPORTED_FEATURES
= (Features
.ALL_PROXY
, Features
.NO_PROXY
)
89 RH_NAME
= 'websockets'
91 def __init__(self
, *args
, **kwargs
):
92 super().__init
__(*args
, **kwargs
)
93 for name
in ('websockets.client', 'websockets.server'):
94 logger
= logging
.getLogger(name
)
95 handler
= logging
.StreamHandler(stream
=sys
.stdout
)
96 handler
.setFormatter(logging
.Formatter(f
'{self.RH_NAME}: %(message)s'))
97 logger
.addHandler(handler
)
99 logger
.setLevel(logging
.DEBUG
)
101 def _check_extensions(self
, extensions
):
102 super()._check
_extensions
(extensions
)
103 extensions
.pop('timeout', None)
104 extensions
.pop('cookiejar', None)
106 def _send(self
, request
):
107 timeout
= float(request
.extensions
.get('timeout') or self
.timeout
)
108 headers
= self
._merge
_headers
(request
.headers
)
109 if 'cookie' not in headers
:
110 cookiejar
= request
.extensions
.get('cookiejar') or self
.cookiejar
111 cookie_header
= cookiejar
.get_cookie_header(request
.url
)
113 headers
['cookie'] = cookie_header
115 wsuri
= parse_uri(request
.url
)
116 create_conn_kwargs
= {
117 'source_address': (self
.source_address
, 0) if self
.source_address
else None,
120 proxy
= select_proxy(request
.url
, request
.proxies
or self
.proxies
or {})
123 socks_proxy_options
= make_socks_proxy_opts(proxy
)
124 sock
= create_connection(
125 address
=(socks_proxy_options
['addr'], socks_proxy_options
['port']),
126 _create_socket_func
=functools
.partial(
127 create_socks_proxy_socket
, (wsuri
.host
, wsuri
.port
), socks_proxy_options
),
131 sock
= create_connection(
132 address
=(wsuri
.host
, wsuri
.port
),
135 conn
= websockets
.sync
.client
.connect(
138 additional_headers
=headers
,
139 open_timeout
=timeout
,
140 user_agent_header
=None,
141 ssl_context
=self
._make
_sslcontext
() if wsuri
.secure
else None,
142 close_timeout
=0, # not ideal, but prevents yt-dlp hanging
144 return WebsocketsResponseAdapter(conn
, url
=request
.url
)
146 # Exceptions as per https://websockets.readthedocs.io/en/stable/reference/sync/client.html
147 except SocksProxyError
as e
:
148 raise ProxyError(cause
=e
) from e
149 except websockets
.exceptions
.InvalidURI
as e
:
150 raise RequestError(cause
=e
) from e
151 except ssl
.SSLCertVerificationError
as e
:
152 raise CertificateVerifyError(cause
=e
) from e
153 except ssl
.SSLError
as e
:
154 raise SSLError(cause
=e
) from e
155 except websockets
.exceptions
.InvalidStatus
as e
:
158 fp
=io
.BytesIO(e
.response
.body
),
160 headers
=e
.response
.headers
,
161 status
=e
.response
.status_code
,
162 reason
=e
.response
.reason_phrase
),
164 except (OSError, TimeoutError
, websockets
.exceptions
.WebSocketException
) as e
:
165 raise TransportError(cause
=e
) from e