ai_wht_wechat/backend/tianqi_proxy_pool.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
大麦IP代理池管理模块
支持动态获取代理IP，实现每次请求使用不同IP
注意：代理用户名和密码从config.dev.yaml的proxy_pool配置中读取
"""
import requests
import time
from typing import Optional, Dict
from loguru import logger
from config import get_config


class TianqiProxyPool:
    """大麦IP代理池管理器"""
    
    def __init__(self):
        """初始化代理池"""
        config = get_config()
        self.enabled = config.get_bool('proxy_pool.enabled', False)
        self.api_url = config.get_str('proxy_pool.api_url', '')
        self.username = config.get_str('proxy_pool.username', '')  # 代理用户名
        self.password = config.get_str('proxy_pool.password', '')  # 代理密码
        self.last_fetch_time = 0
        self.min_fetch_interval = 3  # 最小请求间隔（秒），避免频繁请求API
        
        if self.enabled:
            logger.info(f"[大麦代理池] 已启用，API: {self.api_url[:50]}...")
            if self.username and self.password:
                logger.info(f"[大麦代理池] 使用认证代理，用户名: {self.username}")
            else:
                logger.info(f"[大麦代理池] 使用白名单代理（无认证）")
        else:
            logger.info("[大麦代理池] 未启用")
    
    def is_enabled(self) -> bool:
        """检查代理池是否启用"""
        return self.enabled and bool(self.api_url)
    
    def fetch_proxy(self) -> Optional[Dict[str, str]]:
        """
        从大麦IP API获取一个新的代理IP
        
        Returns:
            代理配置字典 {'server': 'http://ip:port', 'username': '...', 'password': '...'}
            失败返回None
        """
        if not self.is_enabled():
            logger.warning("[大麦代理池] 代理池未启用")
            return None
        
        # 检查请求间隔
        current_time = time.time()
        if current_time - self.last_fetch_time < self.min_fetch_interval:
            wait_time = self.min_fetch_interval - (current_time - self.last_fetch_time)
            logger.info(f"[大麦代理池] 距离上次请求不足{self.min_fetch_interval}秒，等待{wait_time:.1f}秒...")
            time.sleep(wait_time)
        
        try:
            logger.info("[大麦代理池] 正在获取新代理IP...")
            
            # 调用大麦IP API
            response = requests.get(
                self.api_url,
                timeout=10,
                headers={'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'}
            )
            
            self.last_fetch_time = time.time()
            
            if response.status_code != 200:
                logger.error(f"[大麦代理池] API请求失败: HTTP {response.status_code}")
                logger.error(f"[大麦代理池] 返回内容: {response.text[:200]}")
                return None
            
            # 解析返回的IP:Port格式
            proxy_text = response.text.strip()
            
            # 检查是否返回错误信息
            if not proxy_text or ':' not in proxy_text:
                logger.error(f"[大麦代理池] API返回格式错误: {proxy_text}")
                return None
            
            # 构建代理配置（注入用户名密码）
            proxy_config = {
                'server': f'http://{proxy_text}',
                'username': self.username,  # 从配置读取的用户名
                'password': self.password,  # 从配置读取的密码
                'name': f'大麦动态IP-{proxy_text.split(":")[0]}'
            }
            
            if self.username and self.password:
                logger.success(f"[大麦代理池] 获取成功: {proxy_text} (认证代理)")
            else:
                logger.success(f"[大麦代理池] 获取成功: {proxy_text} (白名单代理)")
            return proxy_config
            
        except requests.exceptions.Timeout:
            logger.error("[大麦代理池] API请求超时")
            return None
        except Exception as e:
            logger.error(f"[大麦代理池] 获取代理失败: {str(e)}")
            return None
    
    def format_for_playwright(self, proxy_config: Dict[str, str]) -> Dict[str, str]:
        """
        将代理配置格式化为Playwright格式
        
        Args:
            proxy_config: 代理配置字典
            
        Returns:
            Playwright proxy格式
        """
        result = {
            'server': proxy_config['server']
        }
        
        # 只有在有用户名密码时才添加
        if proxy_config.get('username'):
            result['username'] = proxy_config['username']
        if proxy_config.get('password'):
            result['password'] = proxy_config['password']
        
        return result


# 全局代理池实例
_proxy_pool = None


def get_tianqi_proxy_pool() -> TianqiProxyPool:
    """获取全局大麦代理池实例"""
    global _proxy_pool
    if _proxy_pool is None:
        _proxy_pool = TianqiProxyPool()
    return _proxy_pool


def get_new_proxy() -> Optional[Dict[str, str]]:
    """
    快捷函数：获取一个新的代理IP
    
    Returns:
        Playwright格式的代理配置，失败返回None
    """
    pool = get_tianqi_proxy_pool()
    proxy_config = pool.fetch_proxy()
    
    if proxy_config:
        return pool.format_for_playwright(proxy_config)
    
    return None


if __name__ == "__main__":
    """测试代码"""
    print("=" * 60)
    print("大麦IP代理池测试")
    print("=" * 60)
    
    # 初始化配置
    from config import init_config
    init_config('dev')
    
    # 获取代理池
    pool = get_tianqi_proxy_pool()
    
    if not pool.is_enabled():
        print("❌ 代理池未启用，请在config.dev.yaml中设置 proxy_pool.enabled=true")
    else:
        print("✅ 代理池已启用")
        print(f"   认证模式: {'\u662f' if pool.username and pool.password else '\u5426 (白名单)'}")
        
        # 测试获取3个代理IP
        for i in range(3):
            print(f"\n第{i+1}次获取:")
            proxy = get_new_proxy()
            if proxy:
                print(f"  代理服务器: {proxy['server']}")
                print(f"  有认证: {'是' if proxy.get('username') else '否'}")
                if proxy.get('username'):
                    print(f"  用户名: {proxy['username']}")
            else:
                print("  ❌ 获取失败")
            
            if i < 2:
                print("  等待3秒...")
                time.sleep(3)
    
    print("\n" + "=" * 60)