ai_wht_wechat/backend/test_cookie_inject.py

"""
Cookie注入测试脚本
使用Playwright注入Cookie并验证其有效性
支持跳转到创作者中心或小红书首页
"""
import asyncio
import sys
import json
from pathlib import Path
from playwright.async_api import async_playwright
from typing import Optional, List, Dict, Any


class CookieInjector:
    """Cookie注入器"""
    
    def __init__(self, headless: bool = False):
        """
        初始化Cookie注入器
        
        Args:
            headless: 是否使用无头模式，False可以看到浏览器界面
        """
        self.headless = headless
        self.playwright = None
        self.browser = None
        self.context = None
        self.page = None
    
    async def init_browser(self):
        """初始化浏览器"""
        try:
            print("正在启动浏览器...")
            
            # Windows环境下设置事件循环策略
            if sys.platform == 'win32':
                try:
                    asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
                except Exception as e:
                    print(f"警告: 设置事件循环策略失败: {str(e)}")
            
            self.playwright = await async_playwright().start()
            
            # 启动浏览器
            self.browser = await self.playwright.chromium.launch(
                headless=self.headless,
                args=['--disable-blink-features=AutomationControlled']
            )
            
            # 创建浏览器上下文
            self.context = await self.browser.new_context(
                viewport={'width': 1280, 'height': 720},
                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'
            )
            
            # 创建新页面
            self.page = await self.context.new_page()
            
            print("浏览器初始化成功")
            
        except Exception as e:
            print(f"浏览器初始化失败: {str(e)}")
            raise
    
    async def inject_cookies(self, cookies: List[Dict[str, Any]]) -> bool:
        """
        注入Cookie
        
        Args:
            cookies: Cookie列表
            
        Returns:
            是否注入成功
        """
        try:
            if not self.context:
                await self.init_browser()
            
            print(f"正在注入 {len(cookies)} 个Cookie...")
            
            # 注入Cookie到浏览器上下文
            await self.context.add_cookies(cookies)
            
            print("Cookie注入成功")
            return True
            
        except Exception as e:
            print(f"Cookie注入失败: {str(e)}")
            return False
    
    async def verify_and_navigate(self, target_page: str = 'creator') -> Dict[str, Any]:
        """
        验证Cookie并跳转到指定页面
        
        Args:
            target_page: 目标页面类型 ('creator' 或 'home')
            
        Returns:
            验证结果字典
        """
        try:
            if not self.page:
                return {"success": False, "error": "浏览器未初始化"}
            
            # 确定目标URL
            urls = {
                'creator': 'https://creator.xiaohongshu.com',
                'home': 'https://www.xiaohongshu.com'
            }
            target_url = urls.get(target_page, urls['creator'])
            page_name = '创作者中心' if target_page == 'creator' else '小红书首页'
            
            print(f"\n正在访问{page_name}: {target_url}")
            
            # 访问目标页面
            await self.page.goto(target_url, wait_until='networkidle', timeout=30000)
            await asyncio.sleep(2)  # 等待页面完全加载
            
            # 获取当前URL和标题
            current_url = self.page.url
            title = await self.page.title()
            
            print(f"当前URL: {current_url}")
            print(f"页面标题: {title}")
            
            # 检查是否被重定向到登录页
            is_logged_in = 'login' not in current_url.lower()
            
            if is_logged_in:
                print("Cookie验证成功，已登录状态")
                
                # 尝试获取用户信息
                try:
                    # 等待用户相关元素出现（如头像、用户名等）
                    await self.page.wait_for_selector('[class*="avatar"], [class*="user"]', timeout=5000)
                    print("检测到用户信息元素，确认登录成功")
                except Exception:
                    print("未检测到明显的用户信息元素，但未跳转到登录页")
                
                return {
                    "success": True,
                    "message": f"Cookie有效，已成功访问{page_name}",
                    "url": current_url,
                    "title": title,
                    "logged_in": True
                }
            else:
                print("Cookie可能已失效，页面跳转到登录页")
                return {
                    "success": False,
                    "error": "Cookie已失效或无效，页面跳转到登录页",
                    "url": current_url,
                    "title": title,
                    "logged_in": False
                }
                
        except Exception as e:
            print(f"验证过程异常: {str(e)}")
            import traceback
            traceback.print_exc()
            return {
                "success": False,
                "error": f"验证过程异常: {str(e)}"
            }
    
    async def keep_browser_open(self, duration: int = 60):
        """
        保持浏览器打开一段时间，方便观察
        
        Args:
            duration: 保持打开的秒数，0表示永久打开直到手动关闭
        """
        try:
            if duration == 0:
                print("\n浏览器将保持打开，按 Ctrl+C 关闭...")
                try:
                    while True:
                        await asyncio.sleep(1)
                except KeyboardInterrupt:
                    print("\n用户中断，准备关闭浏览器...")
            else:
                print(f"\n浏览器将保持打开 {duration} 秒...")
                await asyncio.sleep(duration)
                print("时间到，准备关闭浏览器...")
        except Exception as e:
            print(f"保持浏览器异常: {str(e)}")
    
    async def close_browser(self):
        """关闭浏览器"""
        try:
            print("\n正在关闭浏览器...")
            if self.page:
                await self.page.close()
            if self.context:
                await self.context.close()
            if self.browser:
                await self.browser.close()
            if self.playwright:
                await self.playwright.stop()
            print("浏览器已关闭")
        except Exception as e:
            print(f"关闭浏览器异常: {str(e)}")


def load_cookies_from_file(file_path: str) -> Optional[List[Dict[str, Any]]]:
    """
    从文件加载Cookie
    
    Args:
        file_path: Cookie文件路径
        
    Returns:
        Cookie列表，失败返回None
    """
    try:
        cookie_file = Path(file_path)
        if not cookie_file.exists():
            print(f"Cookie文件不存在: {file_path}")
            return None
        
        with open(cookie_file, 'r', encoding='utf-8') as f:
            cookies = json.load(f)
        
        if not isinstance(cookies, list):
            print("Cookie格式错误：必须是数组")
            return None
        
        if len(cookies) == 0:
            print("Cookie数组为空")
            return None
        
        # 验证每个Cookie必须有name和value
        for cookie in cookies:
            if not cookie.get('name') or not cookie.get('value'):
                print(f"Cookie格式错误：缺少name或value字段")
                return None
        
        print(f"成功加载 {len(cookies)} 个Cookie")
        return cookies
        
    except json.JSONDecodeError as e:
        print(f"Cookie文件JSON解析失败: {str(e)}")
        return None
    except Exception as e:
        print(f"加载Cookie文件失败: {str(e)}")
        return None


async def test_cookie_inject(
    cookies_source: str,
    target_page: str = 'creator',
    headless: bool = False,
    keep_open: int = 0
):
    """
    测试Cookie注入
    
    Args:
        cookies_source: Cookie来源（文件路径或JSON字符串）
        target_page: 目标页面 ('creator' 或 'home')
        headless: 是否使用无头模式
        keep_open: 保持浏览器打开的秒数（0表示永久打开）
    """
    print("="*60)
    print("Cookie注入并验证测试")
    print("="*60)
    
    # 加载Cookie
    cookies = None
    
    # 尝试作为文件路径加载
    if Path(cookies_source).exists():
        print(f"\n从文件加载Cookie: {cookies_source}")
        cookies = load_cookies_from_file(cookies_source)
    else:
        # 尝试作为JSON字符串解析
        try:
            print("\n尝试解析Cookie JSON字符串...")
            cookies = json.loads(cookies_source)
            if isinstance(cookies, list) and len(cookies) > 0:
                print(f"成功解析 {len(cookies)} 个Cookie")
        except Exception as e:
            print(f"Cookie解析失败: {str(e)}")
    
    if not cookies:
        print("\n加载Cookie失败，请检查输入")
        return
    
    # 创建注入器
    injector = CookieInjector(headless=headless)
    
    try:
        # 初始化浏览器
        await injector.init_browser()
        
        # 注入Cookie
        inject_success = await injector.inject_cookies(cookies)
        
        if not inject_success:
            print("\nCookie注入失败")
            return
        
        # 验证并跳转
        result = await injector.verify_and_navigate(target_page)
        
        print("\n" + "="*60)
        print("验证结果")
        print("="*60)
        
        if result.get('success'):
            print(f"状态: 成功")
            print(f"消息: {result.get('message')}")
            print(f"URL: {result.get('url')}")
            print(f"标题: {result.get('title')}")
            print(f"登录状态: {'已登录' if result.get('logged_in') else '未登录'}")
        else:
            print(f"状态: 失败")
            print(f"错误: {result.get('error')}")
            if result.get('url'):
                print(f"当前URL: {result.get('url')}")
        
        # 保持浏览器打开
        if keep_open >= 0:
            await injector.keep_browser_open(keep_open)
        
    except KeyboardInterrupt:
        print("\n\n用户中断测试")
    except Exception as e:
        print(f"\n测试过程异常: {str(e)}")
        import traceback
        traceback.print_exc()
    finally:
        await injector.close_browser()
    
    print("\n" + "="*60)
    print("测试完成")
    print("="*60)


async def main():
    """主函数"""
    print("="*60)
    print("小红书Cookie注入测试工具")
    print("="*60)
    
    print("\n功能说明：")
    print("1. 注入Cookie到浏览器")
    print("2. 验证Cookie有效性")
    print("3. 跳转到指定页面（创作者中心/小红书首页）")
    
    print("\n" + "="*60)
    
    # 输入Cookie来源
    print("\n请输入Cookie来源：")
    print("1. 输入Cookie文件路径（如: cookies.json）")
    print("2. 直接粘贴JSON格式的Cookie")
    
    cookies_source = input("\nCookie来源: ").strip()
    
    if not cookies_source:
        print("Cookie来源不能为空")
        return
    
    # 选择目标页面
    print("\n请选择目标页面：")
    print("1. 创作者中心（creator.xiaohongshu.com）")
    print("2. 小红书首页（www.xiaohongshu.com）")
    
    page_choice = input("\n选择 (1 或 2, 默认为 1): ").strip()
    target_page = 'home' if page_choice == '2' else 'creator'
    
    # 选择浏览器模式
    headless_choice = input("\n是否使用无头模式？(y/n, 默认为 n): ").strip().lower()
    headless = headless_choice == 'y'
    
    # 选择保持打开时间
    keep_open_input = input("\n保持浏览器打开时间（秒，0表示直到手动关闭，默认60）: ").strip()
    try:
        keep_open = int(keep_open_input) if keep_open_input else 60
    except ValueError:
        keep_open = 60
    
    # 执行测试
    await test_cookie_inject(
        cookies_source=cookies_source,
        target_page=target_page,
        headless=headless,
        keep_open=keep_open
    )


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