AIstudioProxyAPI / stream /proxy_connector.py
peijun1's picture
Deploy AI Studio Proxy API to Hugging Face Spaces
a5784e9
Raw
History Blame Contribute Delete
2.66 kB
import asyncio
import ssl as ssl_module
import urllib.parse
from typing import Any, Optional, Tuple
from aiohttp import TCPConnector
from python_socks.async_.asyncio import Proxy
class ProxyConnector:
"""
Class to handle connections through different types of proxies
"""
def __init__(self, proxy_url: Optional[str] = None):
self.proxy_url = proxy_url
self.connector = None
if proxy_url:
self._setup_connector()
def _setup_connector(self):
"""Set up the appropriate connector based on the proxy URL"""
if not self.proxy_url:
self.connector = TCPConnector()
return
# Parse the proxy URL
parsed = urllib.parse.urlparse(self.proxy_url)
proxy_type = parsed.scheme.lower()
if proxy_type in ("http", "https", "socks4", "socks5"):
self.connector = "SocksConnector"
else:
raise ValueError(f"Unsupported proxy type: {proxy_type}")
async def create_connection(
self, host: str, port: int, ssl: Optional[Any] = None
) -> Tuple[Any, Any]:
"""Create a connection to the target host through the proxy"""
if not self.connector:
# Direct connection without proxy
reader, writer = await asyncio.open_connection(host, port, ssl=ssl)
return reader, writer
# SOCKS proxy connection
assert self.proxy_url is not None # Type guard for pyright
proxy = Proxy.from_url(self.proxy_url) # type: ignore[misc]
sock = await proxy.connect(dest_host=host, dest_port=port)
if ssl is None:
reader, writer = await asyncio.open_connection(
host=None,
port=None,
sock=sock,
ssl=None,
)
return reader, writer
else:
ssl_context = ssl_module.SSLContext(ssl_module.PROTOCOL_TLS_CLIENT)
ssl_context.check_hostname = False
ssl_context.verify_mode = ssl_module.CERT_NONE
ssl_context.minimum_version = (
ssl_module.TLSVersion.TLSv1_2
) # Force TLS 1.2 or higher
ssl_context.maximum_version = (
ssl_module.TLSVersion.TLSv1_3
) # Allow TLS 1.3 if supported
ssl_context.set_ciphers("DEFAULT@SECLEVEL=2") # Use secure ciphers
reader, writer = await asyncio.open_connection(
host=None,
port=None,
sock=sock,
ssl=ssl_context,
server_hostname=host,
)
return reader, writer