460 lines
14 KiB
Python
460 lines
14 KiB
Python
|
"""
|
|||
|
|
Cookie管理器
|
|||
|
|
统一的Cookie存储、加载、验证和管理功能。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import json
|
|||
|
|
import os
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import Dict, List, Any, Optional
|
|||
|
|
from datetime import datetime, timedelta
|
|||
|
|
from cryptography.fernet import Fernet
|
|||
|
|
import base64
|
|||
|
|
|
|||
|
|
from ..core.models import PlatformType, AccountInfo
|
|||
|
|
from ..config.settings import settings
|
|||
|
|
from ..utils.logger import get_logger
|
|||
|
|
|
|||
|
|
logger = get_logger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class CookieManager:
|
|||
|
|
"""Cookie管理器"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
self.cookie_dir = settings.cookies_dir
|
|||
|
|
self.cookie_dir.mkdir(parents=True, exist_ok=True)
|
|||
|
|
self._encryption_key = None
|
|||
|
|
self._fernet = None
|
|||
|
|
|
|||
|
|
# 初始化加密
|
|||
|
|
if settings.security.encrypt_cookies:
|
|||
|
|
self._init_encryption()
|
|||
|
|
|
|||
|
|
def _init_encryption(self):
|
|||
|
|
"""初始化加密功能"""
|
|||
|
|
try:
|
|||
|
|
key_file = self.cookie_dir / ".encryption_key"
|
|||
|
|
|
|||
|
|
if key_file.exists():
|
|||
|
|
# 读取现有密钥
|
|||
|
|
with open(key_file, 'rb') as f:
|
|||
|
|
self._encryption_key = f.read()
|
|||
|
|
else:
|
|||
|
|
# 生成新密钥
|
|||
|
|
self._encryption_key = Fernet.generate_key()
|
|||
|
|
with open(key_file, 'wb') as f:
|
|||
|
|
f.write(self._encryption_key)
|
|||
|
|
# 设置文件权限为仅所有者可读写
|
|||
|
|
os.chmod(key_file, 0o600)
|
|||
|
|
|
|||
|
|
self._fernet = Fernet(self._encryption_key)
|
|||
|
|
logger.debug("Cookie加密初始化成功")
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Cookie加密初始化失败: {e}")
|
|||
|
|
settings.security.encrypt_cookies = False
|
|||
|
|
|
|||
|
|
def _encrypt_data(self, data: str) -> str:
|
|||
|
|
"""加密数据"""
|
|||
|
|
if not self._fernet:
|
|||
|
|
return data
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
encrypted = self._fernet.encrypt(data.encode('utf-8'))
|
|||
|
|
return base64.b64encode(encrypted).decode('utf-8')
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning(f"数据加密失败: {e}")
|
|||
|
|
return data
|
|||
|
|
|
|||
|
|
def _decrypt_data(self, encrypted_data: str) -> str:
|
|||
|
|
"""解密数据"""
|
|||
|
|
if not self._fernet:
|
|||
|
|
return encrypted_data
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
encrypted_bytes = base64.b64decode(encrypted_data.encode('utf-8'))
|
|||
|
|
decrypted = self._fernet.decrypt(encrypted_bytes)
|
|||
|
|
return decrypted.decode('utf-8')
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning(f"数据解密失败: {e}")
|
|||
|
|
return encrypted_data
|
|||
|
|
|
|||
|
|
def get_cookie_file_path(self, platform: PlatformType, account_name: str) -> Path:
|
|||
|
|
"""
|
|||
|
|
获取Cookie文件路径
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
platform: 平台类型
|
|||
|
|
account_name: 账号名称
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Cookie文件路径
|
|||
|
|
"""
|
|||
|
|
platform_dir = self.cookie_dir / platform.value
|
|||
|
|
platform_dir.mkdir(parents=True, exist_ok=True)
|
|||
|
|
return platform_dir / f"{account_name}.json"
|
|||
|
|
|
|||
|
|
async def save_cookies(
|
|||
|
|
self,
|
|||
|
|
platform: PlatformType,
|
|||
|
|
account_name: str,
|
|||
|
|
cookies: List[Dict[str, Any]],
|
|||
|
|
metadata: Optional[Dict[str, Any]] = None
|
|||
|
|
) -> bool:
|
|||
|
|
"""
|
|||
|
|
保存Cookie
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
platform: 平台类型
|
|||
|
|
account_name: 账号名称
|
|||
|
|
cookies: Cookie列表
|
|||
|
|
metadata: 元数据
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
保存是否成功
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
cookie_file = self.get_cookie_file_path(platform, account_name)
|
|||
|
|
|
|||
|
|
# 准备Cookie数据
|
|||
|
|
cookie_data = {
|
|||
|
|
"cookies": cookies,
|
|||
|
|
"metadata": {
|
|||
|
|
"platform": platform.value,
|
|||
|
|
"account": account_name,
|
|||
|
|
"saved_at": datetime.now().isoformat(),
|
|||
|
|
"expires_at": (datetime.now() + timedelta(days=settings.security.cookie_expiry_days)).isoformat(),
|
|||
|
|
**(metadata or {})
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 序列化为JSON
|
|||
|
|
json_data = json.dumps(cookie_data, indent=2, ensure_ascii=False)
|
|||
|
|
|
|||
|
|
# 加密（如果启用）
|
|||
|
|
if settings.security.encrypt_cookies:
|
|||
|
|
json_data = self._encrypt_data(json_data)
|
|||
|
|
# 保存为加密文件
|
|||
|
|
encrypted_file = cookie_file.with_suffix('.enc')
|
|||
|
|
with open(encrypted_file, 'w', encoding='utf-8') as f:
|
|||
|
|
f.write(json_data)
|
|||
|
|
|
|||
|
|
# 删除未加密文件
|
|||
|
|
if cookie_file.exists():
|
|||
|
|
cookie_file.unlink()
|
|||
|
|
else:
|
|||
|
|
# 直接保存未加密文件
|
|||
|
|
with open(cookie_file, 'w', encoding='utf-8') as f:
|
|||
|
|
f.write(json_data)
|
|||
|
|
|
|||
|
|
logger.info(f"Cookie保存成功: {platform.value}/{account_name}")
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Cookie保存失败: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
async def load_cookies(
|
|||
|
|
self,
|
|||
|
|
platform: PlatformType,
|
|||
|
|
account_name: str
|
|||
|
|
) -> Optional[List[Dict[str, Any]]]:
|
|||
|
|
"""
|
|||
|
|
加载Cookie
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
platform: 平台类型
|
|||
|
|
account_name: 账号名称
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Cookie列表，如果加载失败返回None
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
cookie_file = self.get_cookie_file_path(platform, account_name)
|
|||
|
|
encrypted_file = cookie_file.with_suffix('.enc')
|
|||
|
|
|
|||
|
|
# 确定读取哪个文件
|
|||
|
|
file_to_read = None
|
|||
|
|
if encrypted_file.exists():
|
|||
|
|
file_to_read = encrypted_file
|
|||
|
|
elif cookie_file.exists():
|
|||
|
|
file_to_read = cookie_file
|
|||
|
|
|
|||
|
|
if not file_to_read:
|
|||
|
|
logger.debug(f"Cookie文件不存在: {platform.value}/{account_name}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# 读取文件
|
|||
|
|
with open(file_to_read, 'r', encoding='utf-8') as f:
|
|||
|
|
json_data = f.read()
|
|||
|
|
|
|||
|
|
# 解密（如果是加密文件）
|
|||
|
|
if file_to_read == encrypted_file:
|
|||
|
|
json_data = self._decrypt_data(json_data)
|
|||
|
|
|
|||
|
|
# 解析JSON
|
|||
|
|
cookie_data = json.loads(json_data)
|
|||
|
|
|
|||
|
|
# 验证数据格式
|
|||
|
|
if "cookies" not in cookie_data:
|
|||
|
|
logger.error("Cookie文件格式错误")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
cookies = cookie_data["cookies"]
|
|||
|
|
|
|||
|
|
# 检查是否过期
|
|||
|
|
metadata = cookie_data.get("metadata", {})
|
|||
|
|
expires_at = metadata.get("expires_at")
|
|||
|
|
if expires_at:
|
|||
|
|
try:
|
|||
|
|
expiry_time = datetime.fromisoformat(expires_at.replace('Z', '+00:00'))
|
|||
|
|
if expiry_time < datetime.now():
|
|||
|
|
logger.warning(f"Cookie已过期: {platform.value}/{account_name}")
|
|||
|
|
# 删除过期Cookie
|
|||
|
|
self.delete_cookies(platform, account_name)
|
|||
|
|
return None
|
|||
|
|
except:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
logger.debug(f"Cookie加载成功: {platform.value}/{account_name}")
|
|||
|
|
return cookies
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Cookie加载失败: {e}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def delete_cookies(self, platform: PlatformType, account_name: str) -> bool:
|
|||
|
|
"""
|
|||
|
|
删除Cookie
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
platform: 平台类型
|
|||
|
|
account_name: 账号名称
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
删除是否成功
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
cookie_file = self.get_cookie_file_path(platform, account_name)
|
|||
|
|
encrypted_file = cookie_file.with_suffix('.enc')
|
|||
|
|
|
|||
|
|
deleted = False
|
|||
|
|
if cookie_file.exists():
|
|||
|
|
cookie_file.unlink()
|
|||
|
|
deleted = True
|
|||
|
|
|
|||
|
|
if encrypted_file.exists():
|
|||
|
|
encrypted_file.unlink()
|
|||
|
|
deleted = True
|
|||
|
|
|
|||
|
|
if deleted:
|
|||
|
|
logger.info(f"Cookie删除成功: {platform.value}/{account_name}")
|
|||
|
|
|
|||
|
|
return deleted
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Cookie删除失败: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def list_accounts(self, platform: PlatformType) -> List[str]:
|
|||
|
|
"""
|
|||
|
|
列出指定平台的所有账号
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
platform: 平台类型
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
账号名称列表
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
platform_dir = self.cookie_dir / platform.value
|
|||
|
|
if not platform_dir.exists():
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
accounts = []
|
|||
|
|
for file_path in platform_dir.glob("*.json"):
|
|||
|
|
# 跳过加密密钥文件
|
|||
|
|
if file_path.name.startswith('.'):
|
|||
|
|
continue
|
|||
|
|
accounts.append(file_path.stem)
|
|||
|
|
|
|||
|
|
# 检查加密文件
|
|||
|
|
for file_path in platform_dir.glob("*.enc"):
|
|||
|
|
if file_path.name.startswith('.'):
|
|||
|
|
continue
|
|||
|
|
accounts.append(file_path.stem)
|
|||
|
|
|
|||
|
|
return list(set(accounts)) # 去重
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"获取账号列表失败: {e}")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
async def validate_cookies(
|
|||
|
|
self,
|
|||
|
|
platform: PlatformType,
|
|||
|
|
account_name: str
|
|||
|
|
) -> bool:
|
|||
|
|
"""
|
|||
|
|
验证Cookie是否有效
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
platform: 平台类型
|
|||
|
|
account_name: 账号名称
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Cookie是否有效
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
cookies = await self.load_cookies(platform, account_name)
|
|||
|
|
if not cookies:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
# 基本验证
|
|||
|
|
if not isinstance(cookies, list) or len(cookies) == 0:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
# 检查关键Cookie
|
|||
|
|
important_cookies = self._get_important_cookies(platform)
|
|||
|
|
cookie_names = [c.get('name', '') for c in cookies]
|
|||
|
|
|
|||
|
|
has_important_cookies = any(
|
|||
|
|
name in cookie_names for name in important_cookies
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not has_important_cookies:
|
|||
|
|
logger.warning(f"Cookie缺少关键字段: {platform.value}/{account_name}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Cookie验证失败: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def _get_important_cookies(self, platform: PlatformType) -> List[str]:
|
|||
|
|
"""
|
|||
|
|
获取平台的重要Cookie名称
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
platform: 平台类型
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
重要Cookie名称列表
|
|||
|
|
"""
|
|||
|
|
cookie_mapping = {
|
|||
|
|
PlatformType.XIAOHONGSHU: [
|
|||
|
|
'sessionid', 'token', 'userid', 'webId', 'a1', 'web_session',
|
|||
|
|
'webId', 'web_session', 'x-t', 'webId'
|
|||
|
|
],
|
|||
|
|
PlatformType.DOUYIN: [
|
|||
|
|
'sessionid', 'sid_guard', 'uid_tt', 'sid_tt', 'ttcid',
|
|||
|
|
'passport_csrf_token', '__ac_nonce', '__ac_signature',
|
|||
|
|
'MONITOR_WEB_ID'
|
|||
|
|
]
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return cookie_mapping.get(platform, [])
|
|||
|
|
|
|||
|
|
async def cleanup_expired_cookies(self) -> int:
|
|||
|
|
"""
|
|||
|
|
清理过期的Cookie
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
清理的Cookie数量
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
cleaned_count = 0
|
|||
|
|
|
|||
|
|
for platform in PlatformType:
|
|||
|
|
accounts = self.list_accounts(platform)
|
|||
|
|
for account in accounts:
|
|||
|
|
if not await self.validate_cookies(platform, account):
|
|||
|
|
self.delete_cookies(platform, account)
|
|||
|
|
cleaned_count += 1
|
|||
|
|
logger.info(f"清理过期Cookie: {platform.value}/{account}")
|
|||
|
|
|
|||
|
|
logger.info(f"Cookie清理完成，共清理 {cleaned_count} 个过期Cookie")
|
|||
|
|
return cleaned_count
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Cookie清理失败: {e}")
|
|||
|
|
return 0
|
|||
|
|
|
|||
|
|
async def get_cookie_metadata(
|
|||
|
|
self,
|
|||
|
|
platform: PlatformType,
|
|||
|
|
account_name: str
|
|||
|
|
) -> Optional[Dict[str, Any]]:
|
|||
|
|
"""
|
|||
|
|
获取Cookie元数据
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
platform: 平台类型
|
|||
|
|
account_name: 账号名称
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Cookie元数据
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
cookie_file = self.get_cookie_file_path(platform, account_name)
|
|||
|
|
encrypted_file = cookie_file.with_suffix('.enc')
|
|||
|
|
|
|||
|
|
file_to_read = None
|
|||
|
|
if encrypted_file.exists():
|
|||
|
|
file_to_read = encrypted_file
|
|||
|
|
elif cookie_file.exists():
|
|||
|
|
file_to_read = cookie_file
|
|||
|
|
|
|||
|
|
if not file_to_read:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
with open(file_to_read, 'r', encoding='utf-8') as f:
|
|||
|
|
json_data = f.read()
|
|||
|
|
|
|||
|
|
if file_to_read == encrypted_file:
|
|||
|
|
json_data = self._decrypt_data(json_data)
|
|||
|
|
|
|||
|
|
cookie_data = json.loads(json_data)
|
|||
|
|
return cookie_data.get("metadata", {})
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"获取Cookie元数据失败: {e}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
async def update_cookie_metadata(
|
|||
|
|
self,
|
|||
|
|
platform: PlatformType,
|
|||
|
|
account_name: str,
|
|||
|
|
metadata: Dict[str, Any]
|
|||
|
|
) -> bool:
|
|||
|
|
"""
|
|||
|
|
更新Cookie元数据
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
platform: 平台类型
|
|||
|
|
account_name: 账号名称
|
|||
|
|
metadata: 新的元数据
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
更新是否成功
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
cookies = await self.load_cookies(platform, account_name)
|
|||
|
|
if not cookies:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
current_metadata = await self.get_cookie_metadata(platform, account_name) or {}
|
|||
|
|
updated_metadata = {**current_metadata, **metadata}
|
|||
|
|
updated_metadata['updated_at'] = datetime.now().isoformat()
|
|||
|
|
|
|||
|
|
return await self.save_cookies(platform, account_name, cookies, updated_metadata)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"更新Cookie元数据失败: {e}")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 全局Cookie管理器实例
|
|||
|
|
cookie_manager = CookieManager()
|