warp2protobuf/core/proxy_manager.py

# protobuf2openai/proxy_manager.py
import asyncio
import logging
import os
from datetime import datetime
from typing import Optional

logger = logging.getLogger(__name__)


def _resolve_default_proxy() -> Optional[str]:
    """从环境变量或全局配置解析默认代理地址。"""
    env_proxy = os.getenv("WARP_PROXY_URL") or os.getenv("HTTP_PROXY") or os.getenv("http_proxy")
    if env_proxy:
        env_proxy = env_proxy.strip()
        if env_proxy:
            return env_proxy

    try:
        import config  # type: ignore

        proxy = getattr(config, "PROXY_URL", None)
        if proxy:
            proxy = str(proxy).strip()
            if proxy:
                return proxy
    except Exception:
        pass

    return None


class AsyncProxyManager:
    """异步代理管理器"""

    def __init__(self):
        self.used_identifiers = {}
        self.lock = asyncio.Lock()
        self._default_proxy = _resolve_default_proxy()

    async def cleanup_expired_identifiers(self):
        """清理过期的IP标识"""
        current_time = datetime.now()
        async with self.lock:
            expired_keys = [k for k, v in self.used_identifiers.items() if v < current_time]
            for key in expired_keys:
                del self.used_identifiers[key]

    async def get_proxy(self) -> Optional[str]:
        """获取代理IP，若未配置则返回None表示直连。"""
        return self._default_proxy

    def format_proxy_for_httpx(self, proxy_str: str) -> Optional[str]:
        """格式化代理为httpx可用的格式。"""
        if not proxy_str:
            return None

        proxy_str = proxy_str.strip()
        if not proxy_str:
            return None

        try:
            # 已经是完整URL时直接返回（支持 http/https/socks）
            if proxy_str.startswith(("http://", "https://", "socks5://", "socks4://")):
                return proxy_str

            if '@' in proxy_str:
                credentials, host_port = proxy_str.split('@')
                user, password = credentials.split(':')
                host, port = host_port.split(':')
                return f"socks5://{user}:{password}@{host}:{port}"

            # host:port 形式，默认按 socks5 处理
            if ':' in proxy_str:
                host, port = proxy_str.split(':', 1)
                return f"socks5://{host}:{port}"

            logger.error(f"代理格式无法识别: {proxy_str}")
            return None
        except Exception as e:
            logger.error(f"格式化代理失败: {e}")
            return None