ai_wht_wechat/backend/test_login_flow.py

"""
小红书验证码登录流程测试脚本
测试完整的验证码发送和登录流程
"""
import asyncio
import sys
from xhs_login import XHSLoginService


async def test_send_verification_code(phone: str, proxy_index: int = 0):
    """
    测试发送验证码流程
    
    Args:
        phone: 手机号
        proxy_index: 代理索引 (0 或 1)
    """
    print(f"\n{'='*60}")
    print(f"📱 测试发送验证码流程")
    print(f"{'='*60}")
    
    # 从代理配置获取代理信息
    from damai_proxy_config import get_proxy_config
    proxy_config = get_proxy_config(proxy_index)
    proxy_server = proxy_config['server'].replace('http://', '')
    proxy_url = f"http://{proxy_config['username']}:{proxy_config['password']}@{proxy_server}"
    
    print(f"✅ 使用代理: 代理{proxy_index + 1}")
    print(f"   代理服务器: {proxy_config['server']}")
    print(f"   手机号: {phone}")
    
    # 创建登录服务
    login_service = XHSLoginService()
    
    try:
        # 初始化浏览器（使用代理）
        user_agent = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
        await login_service.init_browser(proxy=proxy_url, user_agent=user_agent)
        print("✅ 浏览器初始化成功（已启用代理）")
        
        # 发送验证码
        print(f"\n📤 正在发送验证码到 {phone}...")
        result = await login_service.send_verification_code(phone)
        
        if result.get('success'):
            print(f"✅ 验证码发送成功！")
            print(f"   消息: {result.get('message')}")
            return login_service  # 返回服务实例供后续登录使用
        else:
            print(f"❌ 验证码发送失败: {result.get('error')}")
            return None
            
    except Exception as e:
        print(f"❌ 发送验证码过程异常: {str(e)}")
        import traceback
        traceback.print_exc()
        return None


async def test_login_with_code(login_service: XHSLoginService, phone: str, code: str):
    """
    测试使用验证码登录
    
    Args:
        login_service: XHSLoginService实例
        phone: 手机号
        code: 验证码
    """
    print(f"\n{'='*60}")
    print(f"🔑 测试使用验证码登录")
    print(f"{'='*60}")
    
    print(f"   手机号: {phone}")
    print(f"   验证码: {code}")
    
    try:
        # 执行登录
        result = await login_service.login(phone, code)
        
        if result.get('success'):
            print("✅ 登录成功！")
            
            # 显示获取到的Cookies信息
            cookies = result.get('cookies', {})
            print(f"   获取到 {len(cookies)} 个Cookie")
            
            # 保存完整Cookies到文件
            cookies_full = result.get('cookies_full', [])
            if cookies_full:
                import json
                with open('cookies.json', 'w', encoding='utf-8') as f:
                    json.dump(cookies_full, f, ensure_ascii=False, indent=2)
                print("   ✅ 已保存完整Cookies到 cookies.json")
            
            # 显示用户信息
            user_info = result.get('user_info', {})
            if user_info:
                print(f"   用户信息: {list(user_info.keys())}")
            
            return result
        else:
            print(f"❌ 登录失败: {result.get('error')}")
            return result
            
    except Exception as e:
        print(f"❌ 登录过程异常: {str(e)}")
        import traceback
        traceback.print_exc()
        return {"success": False, "error": str(e)}


async def test_complete_login_flow(phone: str, code: str = None, proxy_index: int = 0):
    """
    测试完整的登录流程
    
    Args:
        phone: 手机号
        code: 验证码（如果为None，则只测试发送验证码）
        proxy_index: 代理索引
    """
    print("="*60)
    print("🔄 测试完整登录流程")
    print("="*60)
    
    # 步骤1: 发送验证码
    print("\n📋 步骤1: 发送验证码")
    login_service = await test_send_verification_code(phone, proxy_index)
    
    if not login_service:
        print("❌ 发送验证码失败，终止流程")
        return
    
    # 如果提供了验证码，则执行登录
    if code:
        print("\n📋 步骤2: 使用验证码登录")
        result = await test_login_with_code(login_service, phone, code)
        
        if result.get('success'):
            print("\n🎉 完整登录流程成功！")
        else:
            print(f"\n❌ 完整登录流程失败: {result.get('error')}")
    else:
        print("\n⚠️  提供了验证码参数才可完成登录步骤")
        print("   请在手机上查看验证码，然后调用登录方法")
    
    # 清理资源
    await login_service.close_browser()


async def test_multiple_proxies_login(phone: str, proxy_indices: list = [0, 1]):
    """
    测试使用多个代理进行登录
    
    Args:
        phone: 手机号
        proxy_indices: 代理索引列表
    """
    print("="*60)
    print("🔄 测试多代理登录")
    print("="*60)
    
    for i, proxy_idx in enumerate(proxy_indices):
        print(f"\n🧪 测试代理 {proxy_idx + 1} (第 {i+1} 次尝试)")
        
        # 由于验证码只能发送一次，这里只测试发送验证码
        login_service = await test_send_verification_code(phone, proxy_idx)
        
        if login_service:
            print(f"   ✅ 代理 {proxy_idx + 1} 发送验证码成功")
            await login_service.close_browser()
        else:
            print(f"   ❌ 代理 {proxy_idx + 1} 发送验证码失败")
        
        # 在测试之间添加延迟
        if i < len(proxy_indices) - 1:
            print("   ⏳ 等待3秒后测试下一个代理...")
            await asyncio.sleep(3)


def show_usage_examples():
    """显示使用示例"""
    print("="*60)
    print("💡 使用示例")
    print("="*60)
    
    print("\n1️⃣ 仅发送验证码:")
    print("   # 发送验证码到手机号，使用代理1")
    print("   await test_send_verification_code('13800138000', proxy_index=0)")
    
    print("\n2️⃣ 完整登录流程:")
    print("   # 完整流程：发送验证码 + 登录")
    print("   await test_complete_login_flow('13800138000', '123456', proxy_index=0)")
    
    print("\n3️⃣ 多代理测试:")
    print("   # 测试多个代理")
    print("   await test_multiple_proxies_login('13800138000', [0, 1])")


async def main():
    """主函数"""
    show_usage_examples()
    
    print(f"\n{'='*60}")
    print("🎯 选择测试模式")
    print("="*60)
    
    print("\n1. 发送验证码测试")
    print("2. 完整登录流程测试")
    print("3. 多代理测试")
    
    try:
        choice = input("\n请选择测试模式 (1-3, 默认为1): ").strip()
        
        if choice not in ['1', '2', '3']:
            choice = '1'
        
        phone = input("请输入手机号: ").strip()
        
        if not phone:
            print("❌ 手机号不能为空")
            return
        
        if choice == '1':
            proxy_choice = input("请选择代理 (0 或 1, 默认为0): ").strip()
            if proxy_choice not in ['0', '1']:
                proxy_choice = '0'
            proxy_idx = int(proxy_choice)
            
            await test_send_verification_code(phone, proxy_idx)
            
        elif choice == '2':
            code = input("请输入验证码 (留空则只测试发送): ").strip()
            proxy_choice = input("请选择代理 (0 或 1, 默认为0): ").strip()
            if proxy_choice not in ['0', '1']:
                proxy_choice = '0'
            proxy_idx = int(proxy_choice)
            
            await test_complete_login_flow(phone, code if code else None, proxy_idx)
            
        elif choice == '3':
            await test_multiple_proxies_login(phone)
        
        print(f"\n{'='*60}")
        print("✅ 测试完成！")
        print("="*60)
        
    except KeyboardInterrupt:
        print("\n\n⚠️  测试被用户中断")
    except Exception as e:
        print(f"\n❌ 测试过程中出现错误: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    # Windows环境下设置事件循环策略
    if sys.platform == 'win32':
        asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
    
    # 运行测试
    asyncio.run(main())