""" 统一的Cookie管理器 整合JSON文件和环境变量cookie的检测、加载和管理功能 """ import os import json from dataclasses import dataclass from typing import List, Dict, Optional from utils.paths import cookies_dir from utils.cookie_handler import convert_cookie_editor_to_playwright, convert_kv_to_playwright from utils.common import clean_env_value @dataclass class CookieSource: """Cookie来源的统一表示""" type: str # "file" | "env_var" identifier: str # filename or "USER_COOKIE_1" display_name: str # 显示名称 exists: bool = True def __str__(self): return f"{self.type}:{self.identifier}" class CookieManager: """ 统一的Cookie管理器 负责检测、加载和缓存所有来源的cookie数据 """ def __init__(self, logger=None): self.logger = logger self._detected_sources: Optional[List[CookieSource]] = None self._cookie_cache: Dict[str, List[Dict]] = {} def detect_all_sources(self) -> List[CookieSource]: """ 检测所有可用的cookie来源(JSON文件 + 环境变量) 结果会被缓存,避免重复扫描 """ if self._detected_sources is not None: return self._detected_sources sources = [] # 1. 扫描cookies目录中的JSON文件 try: cookie_path = cookies_dir() if os.path.isdir(cookie_path): cookie_files = [f for f in os.listdir(cookie_path) if f.lower().endswith('.json')] for cookie_file in cookie_files: source = CookieSource( type="file", identifier=cookie_file, display_name=cookie_file ) sources.append(source) if cookie_files and self.logger: self.logger.info(f"发现 {len(cookie_files)} 个 Cookie 文件") elif self.logger: self.logger.info(f"在 {cookie_path} 目录下未找到任何 .json 格式的 Cookie 文件") else: if self.logger: self.logger.error(f"Cookie 目录不存在: {cookie_path}") except Exception as e: if self.logger: self.logger.error(f"扫描 Cookie 目录时出错: {e}") # 2. 扫描USER_COOKIE环境变量 cookie_index = 1 env_cookie_count = 0 while True: env_var_name = f"USER_COOKIE_{cookie_index}" env_value = clean_env_value(os.getenv(env_var_name)) if not env_value: if cookie_index == 1 and self.logger: self.logger.info(f"未检测到任何 USER_COOKIE 环境变量") break source = CookieSource( type="env_var", identifier=env_var_name, display_name=env_var_name ) sources.append(source) env_cookie_count += 1 cookie_index += 1 if env_cookie_count > 0 and self.logger: self.logger.info(f"发现 {env_cookie_count} 个 Cookie 环境变量") # 缓存结果 self._detected_sources = sources return sources def load_cookies(self, source: CookieSource) -> List[Dict]: """ 从指定来源加载cookie数据 Args: source: Cookie来源对象 Returns: Playwright兼容的cookie列表 """ cache_key = str(source) # 检查缓存 if cache_key in self._cookie_cache: if self.logger: self.logger.debug(f"从缓存加载 Cookie: {source.display_name}") return self._cookie_cache[cache_key] cookies = [] try: if source.type == "file": cookies = self._load_from_file(source.identifier) elif source.type == "env_var": cookies = self._load_from_env(source.identifier) else: if self.logger: self.logger.error(f"未知的 Cookie 来源类型: {source.type}") return [] # 缓存结果 self._cookie_cache[cache_key] = cookies if self.logger: self.logger.info(f"从 {source.display_name} 加载了 {len(cookies)} 个 Cookie 数据") except Exception as e: if self.logger: self.logger.error(f"从 {source.display_name} 加载 Cookie 时出错: {e}") return [] return cookies def _load_from_file(self, filename: str) -> List[Dict]: """从JSON文件加载 Cookie""" cookie_path = cookies_dir() / filename if not os.path.exists(cookie_path): raise FileNotFoundError(f"Cookie 文件不存在: {cookie_path}") with open(cookie_path, 'r', encoding='utf-8') as f: cookies_from_file = json.load(f) return convert_cookie_editor_to_playwright(cookies_from_file, logger=self.logger) def _load_from_env(self, env_var_name: str) -> List[Dict]: """从环境变量加载 Cookie""" env_value = clean_env_value(os.getenv(env_var_name)) if not env_value: raise ValueError(f"环境变量 {env_var_name} 不存在或为空") return convert_kv_to_playwright( env_value, default_domain=".google.com", logger=self.logger ) def get_all_sources(self) -> List[CookieSource]: """获取所有检测到的 Cookie 来源""" return self.detect_all_sources() def clear_cache(self): """清空 Cookie 缓存""" self._cookie_cache.clear() if self.logger: self.logger.debug("Cookie 缓存已清空") def get_source_summary(self) -> Dict[str, int]: """ 获取 Cookie 来源统计信息 Returns: 包含各类型来源数量的字典 """ sources = self.detect_all_sources() summary = { "total": len(sources), "files": 0, "env_vars": 0 } for source in sources: if source.type == "file": summary["files"] += 1 elif source.type == "env_var": summary["env_vars"] += 1 return summary