Working websockets HTTP/S proxy

This commit is contained in:
coletdjnz 2024-04-06 15:14:59 +13:00
parent fddf9e0577
commit 3999a510f7
No known key found for this signature in database
GPG Key ID: 91984263BB39894A
3 changed files with 60 additions and 18 deletions

View File

@ -116,6 +116,9 @@ def _io_refs(self):
@_io_refs.setter @_io_refs.setter
def _io_refs(self, value): def _io_refs(self, value):
self.socket._io_refs = value self.socket._io_refs = value
def shutdown(self, *args, **kwargs):
self.socket.shutdown(*args, **kwargs)
else: else:
SSLTransport = None SSLTransport = None
@ -142,13 +145,14 @@ def handle(self):
protocol = websockets.ServerProtocol() protocol = websockets.ServerProtocol()
connection = websockets.sync.server.ServerConnection(socket=self.request, protocol=protocol, close_timeout=0) connection = websockets.sync.server.ServerConnection(socket=self.request, protocol=protocol, close_timeout=0)
connection.handshake() connection.handshake()
for message in connection:
if message == 'proxy_info':
connection.send(json.dumps(self.proxy_info)) connection.send(json.dumps(self.proxy_info))
connection.close() connection.close()
class WebSocketSecureProxyHandler(WebSocketProxyHandler): class WebSocketSecureProxyHandler(WebSocketProxyHandler):
def __init__(self, request, *args, proxy_info=None, **kwargs): def __init__(self, request, *args, **kwargs):
self.proxy_info = proxy_info
certfn = os.path.join(TEST_DIR, 'testcert.pem') certfn = os.path.join(TEST_DIR, 'testcert.pem')
sslctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) sslctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
sslctx.load_cert_chain(certfn, None) sslctx.load_cert_chain(certfn, None)
@ -218,7 +222,7 @@ def proxy_server(proxy_server_class, request_handler, bind_ip=None, **proxy_serv
finally: finally:
server.shutdown() server.shutdown()
server.server_close() server.server_close()
server_thread.join(2.0) server_thread.join()
class HTTPProxyTestContext(abc.ABC): class HTTPProxyTestContext(abc.ABC):
@ -297,6 +301,7 @@ def test_http_no_auth(self, handler, ctx):
proxy_info = ctx.proxy_info_request(rh) proxy_info = ctx.proxy_info_request(rh)
assert proxy_info['connect'] is False assert proxy_info['connect'] is False
assert 'Proxy-Authorization' not in proxy_info['headers'] assert 'Proxy-Authorization' not in proxy_info['headers']
assert proxy_info['proxy'] == server_address
def test_http_auth(self, handler, ctx): def test_http_auth(self, handler, ctx):
with ctx.http_server(HTTPProxyHandler, username='test', password='test') as server_address: with ctx.http_server(HTTPProxyHandler, username='test', password='test') as server_address:
@ -318,8 +323,9 @@ def test_http_source_address(self, handler, ctx):
verify_address_availability(source_address) verify_address_availability(source_address)
with handler(proxies={ctx.REQUEST_PROTO: f'http://{server_address}'}, with handler(proxies={ctx.REQUEST_PROTO: f'http://{server_address}'},
source_address=source_address) as rh: source_address=source_address) as rh:
response = ctx.proxy_info_request(rh) proxy_info = ctx.proxy_info_request(rh)
assert response['client_address'][0] == source_address assert proxy_info['client_address'][0] == source_address
assert proxy_info['proxy'] == server_address
@pytest.mark.skip_handler('Urllib', 'urllib does not support https proxies') @pytest.mark.skip_handler('Urllib', 'urllib does not support https proxies')
def test_https(self, handler, ctx): def test_https(self, handler, ctx):
@ -328,6 +334,7 @@ def test_https(self, handler, ctx):
proxy_info = ctx.proxy_info_request(rh) proxy_info = ctx.proxy_info_request(rh)
assert proxy_info['connect'] is False assert proxy_info['connect'] is False
assert 'Proxy-Authorization' not in proxy_info['headers'] assert 'Proxy-Authorization' not in proxy_info['headers']
assert proxy_info['proxy'] == server_address
@pytest.mark.skip_handler('Urllib', 'urllib does not support https proxies') @pytest.mark.skip_handler('Urllib', 'urllib does not support https proxies')
def test_https_verify_failed(self, handler, ctx): def test_https_verify_failed(self, handler, ctx):
@ -345,6 +352,7 @@ def test_http_with_idn(self, handler, ctx):
proxy_info = ctx.proxy_info_request(rh, target_domain='中文.tw') proxy_info = ctx.proxy_info_request(rh, target_domain='中文.tw')
assert proxy_info['path'].startswith('http://xn--fiq228c.tw') assert proxy_info['path'].startswith('http://xn--fiq228c.tw')
assert proxy_info['headers']['Host'].split(':', 1)[0] == 'xn--fiq228c.tw' assert proxy_info['headers']['Host'].split(':', 1)[0] == 'xn--fiq228c.tw'
assert proxy_info['proxy'] == server_address
@pytest.mark.parametrize( @pytest.mark.parametrize(
@ -361,12 +369,14 @@ def test_http_connect_no_auth(self, handler, ctx):
proxy_info = ctx.proxy_info_request(rh) proxy_info = ctx.proxy_info_request(rh)
assert proxy_info['connect'] is True assert proxy_info['connect'] is True
assert 'Proxy-Authorization' not in proxy_info['headers'] assert 'Proxy-Authorization' not in proxy_info['headers']
assert proxy_info['proxy'] == server_address
def test_http_connect_auth(self, handler, ctx): def test_http_connect_auth(self, handler, ctx):
with ctx.http_server(HTTPConnectProxyHandler, username='test', password='test') as server_address: with ctx.http_server(HTTPConnectProxyHandler, username='test', password='test') as server_address:
with handler(verify=False, proxies={ctx.REQUEST_PROTO: f'http://test:test@{server_address}'}) as rh: with handler(verify=False, proxies={ctx.REQUEST_PROTO: f'http://test:test@{server_address}'}) as rh:
proxy_info = ctx.proxy_info_request(rh) proxy_info = ctx.proxy_info_request(rh)
assert 'Proxy-Authorization' in proxy_info['headers'] assert 'Proxy-Authorization' in proxy_info['headers']
assert proxy_info['proxy'] == server_address
def test_http_connect_bad_auth(self, handler, ctx): def test_http_connect_bad_auth(self, handler, ctx):
with ctx.http_server(HTTPConnectProxyHandler, username='test', password='test') as server_address: with ctx.http_server(HTTPConnectProxyHandler, username='test', password='test') as server_address:
@ -381,8 +391,9 @@ def test_http_connect_source_address(self, handler, ctx):
with handler(proxies={ctx.REQUEST_PROTO: f'http://{server_address}'}, with handler(proxies={ctx.REQUEST_PROTO: f'http://{server_address}'},
source_address=source_address, source_address=source_address,
verify=False) as rh: verify=False) as rh:
response = ctx.proxy_info_request(rh) proxy_info = ctx.proxy_info_request(rh)
assert response['client_address'][0] == source_address assert proxy_info['client_address'][0] == source_address
assert proxy_info['proxy'] == server_address
@pytest.mark.skipif(urllib3 is None, reason='requires urllib3 to test') @pytest.mark.skipif(urllib3 is None, reason='requires urllib3 to test')
def test_https_connect_proxy(self, handler, ctx): def test_https_connect_proxy(self, handler, ctx):
@ -391,6 +402,7 @@ def test_https_connect_proxy(self, handler, ctx):
proxy_info = ctx.proxy_info_request(rh) proxy_info = ctx.proxy_info_request(rh)
assert proxy_info['connect'] is True assert proxy_info['connect'] is True
assert 'Proxy-Authorization' not in proxy_info['headers'] assert 'Proxy-Authorization' not in proxy_info['headers']
assert proxy_info['proxy'] == server_address
@pytest.mark.skipif(urllib3 is None, reason='requires urllib3 to test') @pytest.mark.skipif(urllib3 is None, reason='requires urllib3 to test')
def test_https_connect_verify_failed(self, handler, ctx): def test_https_connect_verify_failed(self, handler, ctx):
@ -408,3 +420,4 @@ def test_https_connect_proxy_auth(self, handler, ctx):
with handler(verify=False, proxies={ctx.REQUEST_PROTO: f'https://test:test@{server_address}'}) as rh: with handler(verify=False, proxies={ctx.REQUEST_PROTO: f'https://test:test@{server_address}'}) as rh:
proxy_info = ctx.proxy_info_request(rh) proxy_info = ctx.proxy_info_request(rh)
assert 'Proxy-Authorization' in proxy_info['headers'] assert 'Proxy-Authorization' in proxy_info['headers']
assert proxy_info['proxy'] == server_address

View File

@ -216,6 +216,8 @@ def handle(self):
protocol = websockets.ServerProtocol() protocol = websockets.ServerProtocol()
connection = websockets.sync.server.ServerConnection(socket=self.request, protocol=protocol, close_timeout=0) connection = websockets.sync.server.ServerConnection(socket=self.request, protocol=protocol, close_timeout=0)
connection.handshake() connection.handshake()
for message in connection:
if message == 'socks_info':
connection.send(json.dumps(self.socks_info)) connection.send(json.dumps(self.socks_info))
connection.close() connection.close()

View File

@ -118,7 +118,7 @@ def __init__(self, *args, **kwargs):
for name in ('websockets.client', 'websockets.server'): for name in ('websockets.client', 'websockets.server'):
logger = logging.getLogger(name) logger = logging.getLogger(name)
handler = logging.StreamHandler(stream=sys.stdout) handler = logging.StreamHandler(stream=sys.stdout)
handler.setFormatter(logging.Formatter(f'{self.RH_NAME}: %(message)s')) handler.setFormatter(logging.Formatter(f'{self.RH_NAME}: [{name}] %(message)s'))
self.__logging_handlers[name] = handler self.__logging_handlers[name] = handler
logger.addHandler(handler) logger.addHandler(handler)
if self.verbose: if self.verbose:
@ -152,7 +152,7 @@ def _make_sock(self, proxy, url, timeout):
**create_conn_kwargs **create_conn_kwargs
) )
elif parsed_proxy_url.scheme.startswith('http'): elif parsed_proxy_url.scheme in ('http', 'https'):
return create_http_connect_conn( return create_http_connect_conn(
proxy_url=proxy, proxy_url=proxy,
url=url, url=url,
@ -177,6 +177,7 @@ def _send(self, request):
headers['cookie'] = cookie_header headers['cookie'] = cookie_header
proxy = select_proxy(request.url, self._get_proxies(request)) proxy = select_proxy(request.url, self._get_proxies(request))
try: try:
conn = websockets.sync.client.connect( conn = websockets.sync.client.connect(
sock=self._make_sock(proxy, request.url, timeout), sock=self._make_sock(proxy, request.url, timeout),
@ -184,7 +185,10 @@ def _send(self, request):
additional_headers=headers, additional_headers=headers,
open_timeout=timeout, open_timeout=timeout,
user_agent_header=None, user_agent_header=None,
ssl_context=self._make_sslcontext() if parse_uri(request.url).secure else None, ssl_context=(
WebsocketsSSLContext(self._make_sslcontext())
if parse_uri(request.url).secure else None
),
close_timeout=0, # not ideal, but prevents yt-dlp hanging close_timeout=0, # not ideal, but prevents yt-dlp hanging
) )
return WebsocketsResponseAdapter(conn, url=request.url) return WebsocketsResponseAdapter(conn, url=request.url)
@ -218,12 +222,34 @@ def begin(self):
if not self._check_close() and not self.chunked and self.length is None: if not self._check_close() and not self.chunked and self.length is None:
self.will_close = False self.will_close = False
class CustomSSLTransport(SSLTransport):
# todo: only define if urllib3 is available
class WebsocketsSSLTransport(SSLTransport):
"""
Modified version of urllib3 SSLTransport to support additional operations used by websockets
"""
def setsockopt(self, *args, **kwargs): def setsockopt(self, *args, **kwargs):
self.socket.setsockopt(*args, **kwargs) self.socket.setsockopt(*args, **kwargs)
def shutdown(self, *args, **kwargs): def shutdown(self, *args, **kwargs):
self.unwrap()
self.socket.shutdown(*args, **kwargs) self.socket.shutdown(*args, **kwargs)
class WebsocketsSSLContext:
"""
Dummy SSL Context for websockets which returns a WebsocketsSSLTransport instance
for wrap socket when using TLS-in-TLS.
"""
def __init__(self, ssl_context: ssl.SSLContext):
self.ssl_context = ssl_context
def wrap_socket(self, sock, server_hostname=None):
if isinstance(sock, ssl.SSLSocket):
return WebsocketsSSLTransport(sock, self.ssl_context, server_hostname=server_hostname)
return self.ssl_context.wrap_socket(sock, server_hostname=server_hostname)
def create_http_connect_conn( def create_http_connect_conn(
proxy_url, proxy_url,
url, url,
@ -256,17 +282,18 @@ def create_http_connect_conn(
if source_address is not None: if source_address is not None:
conn.source_address = (source_address, 0) conn.source_address = (source_address, 0)
conn.debuglevel=2
try: try:
conn.connect() conn.connect()
if ssl_context: if ssl_context:
conn.sock = CustomSSLTransport(conn.sock, ssl_context, server_hostname=proxy_url_parsed.hostname) conn.sock = ssl_context.wrap_socket(conn.sock, server_hostname=proxy_url_parsed.hostname)
conn.request(
conn.request(method='CONNECT', url=f'{request_url_parsed.host}:{request_url_parsed.port}', headers=proxy_headers) method='CONNECT',
url=f'{request_url_parsed.host}:{request_url_parsed.port}',
headers=proxy_headers)
response = conn.getresponse() response = conn.getresponse()
except OSError as e: except OSError as e:
conn.close() conn.close()
raise TransportError('Unable to connect to proxy', cause=e) from e raise ProxyError('Unable to connect to proxy', cause=e) from e
if response.status == 200: if response.status == 200:
return conn.sock return conn.sock