"""
Cookie 管理模块
负责 cookies 的加载、保存和删除
"""
import json
import os
from abc import ABC, abstractmethod
from pathlib import Path
from typing import List, Dict, Any, Optional
from loguru import logger
class Cookier(ABC):
"""Cookie 管理接口"""
@abstractmethod
async def load_cookies(self) -> List[Dict[str, Any]]:
"""加载 cookies"""
pass
@abstractmethod
async def save_cookies(self, cookies: List[Dict[str, Any]]) -> None:
"""保存 cookies"""
pass
@abstractmethod
async def delete_cookies(self) -> None:
"""删除 cookies"""
pass
class LocalCookie(Cookier):
"""本地文件 Cookie 管理"""
def __init__(self, path: str):
if not path:
raise ValueError("path is required")
self.path = Path(path)
async def load_cookies(self) -> List[Dict[str, Any]]:
"""从文件中加载 cookies"""
try:
if not self.path.exists():
logger.warning(f"Cookies file not found: {self.path}")
return []
with open(self.path, "r", encoding="utf-8") as f:
data = json.load(f)
logger.info(f"Loaded {len(data)} cookies from {self.path}")
return data
except Exception as e:
logger.error(f"Failed to load cookies from {self.path}: {e}")
return []
async def save_cookies(self, cookies: List[Dict[str, Any]]) -> None:
"""保存 cookies 到文件"""
try:
# 确保目录存在
self.path.parent.mkdir(parents=True, exist_ok=True)
with open(self.path, "w", encoding="utf-8") as f:
json.dump(cookies, f, indent=2, ensure_ascii=False)
logger.info(f"Saved {len(cookies)} cookies to {self.path}")
except Exception as e:
logger.error(f"Failed to save cookies to {self.path}: {e}")
raise
async def delete_cookies(self) -> None:
"""删除 cookies 文件"""
try:
if self.path.exists():
self.path.unlink()
logger.info(f"Deleted cookies file: {self.path}")
else:
logger.info(f"Cookies file not found, nothing to delete: {self.path}")
except Exception as e:
logger.error(f"Failed to delete cookies file {self.path}: {e}")
raise
def get_cookies_file_path() -> str:
"""
获取 cookies 文件路径
为了向后兼容，如果旧路径 /tmp/cookies.json 存在，则继续使用；
否则使用环境变量 COOKIES_PATH 或当前目录下的 cookies.json
"""
import tempfile
# 旧路径：/tmp/cookies.json
old_path = Path(tempfile.gettempdir()) / "cookies.json"
# 检查旧路径文件是否存在
if old_path.exists():
logger.info(f"Using legacy cookies path: {old_path}")
return str(old_path)
# 使用环境变量或默认路径
path = os.getenv("COOKIES_PATH", "cookies.json")
logger.info(f"Using cookies path: {path}")
return path
def create_cookie_manager(path: Optional[str] = None) -> Cookier:
"""创建 Cookie 管理器"""
if path is None:
path = get_cookies_file_path()
return LocalCookie(path)
