251 lines
9.7 KiB
Python
251 lines
9.7 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
|
|||
|
|
"""
|
|||
|
|
Cookie管理器
|
|||
|
|
支持多个cookies的管理、验证和轮换使用
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import json
|
|||
|
|
import os
|
|||
|
|
import time
|
|||
|
|
import random
|
|||
|
|
from typing import Dict, List, Optional, Tuple
|
|||
|
|
from datetime import datetime, timedelta
|
|||
|
|
import logging
|
|||
|
|
|
|||
|
|
from .models import CookieInfo, CookieStats
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
class CookieManager:
|
|||
|
|
"""Cookie管理器"""
|
|||
|
|
|
|||
|
|
def __init__(self, config_path: str = "config/cookies.json"):
|
|||
|
|
self.config_path = config_path
|
|||
|
|
self.cookies: Dict[str, CookieInfo] = {}
|
|||
|
|
self.current_cookie: Optional[str] = None
|
|||
|
|
self.max_failures = 3 # 最大失败次数
|
|||
|
|
self.rotation_interval = 300 # 轮换间隔（秒）
|
|||
|
|
self.load_cookies()
|
|||
|
|
|
|||
|
|
def load_cookies(self) -> None:
|
|||
|
|
"""加载cookies配置"""
|
|||
|
|
try:
|
|||
|
|
if os.path.exists(self.config_path):
|
|||
|
|
with open(self.config_path, 'r', encoding='utf-8') as f:
|
|||
|
|
data = json.load(f)
|
|||
|
|
|
|||
|
|
for cookie_data in data.get('cookies', []):
|
|||
|
|
cookie_info = CookieInfo(
|
|||
|
|
name=cookie_data['name'],
|
|||
|
|
cookie_string=cookie_data['cookie_string'],
|
|||
|
|
last_used=cookie_data.get('last_used', datetime.now().isoformat()),
|
|||
|
|
use_count=cookie_data.get('use_count', 0),
|
|||
|
|
is_valid=cookie_data.get('is_valid', True),
|
|||
|
|
failure_count=cookie_data.get('failure_count', 0),
|
|||
|
|
user_info=cookie_data.get('user_info')
|
|||
|
|
)
|
|||
|
|
self.cookies[cookie_info.name] = cookie_info
|
|||
|
|
|
|||
|
|
logger.info(f"加载了 {len(self.cookies)} 个cookies")
|
|||
|
|
else:
|
|||
|
|
logger.warning(f"Cookie配置文件不存在: {self.config_path}")
|
|||
|
|
self._create_default_config()
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"加载cookies失败: {str(e)}")
|
|||
|
|
self._create_default_config()
|
|||
|
|
|
|||
|
|
def _create_default_config(self) -> None:
|
|||
|
|
"""创建默认配置文件"""
|
|||
|
|
os.makedirs(os.path.dirname(self.config_path), exist_ok=True)
|
|||
|
|
default_config = {
|
|||
|
|
"cookies": [
|
|||
|
|
{
|
|||
|
|
"name": "default",
|
|||
|
|
"cookie_string": "请在此处填入您的小红书cookie",
|
|||
|
|
"last_used": datetime.now().isoformat(),
|
|||
|
|
"use_count": 0,
|
|||
|
|
"is_valid": True,
|
|||
|
|
"failure_count": 0,
|
|||
|
|
"user_info": None
|
|||
|
|
}
|
|||
|
|
]
|
|||
|
|
}
|
|||
|
|
with open(self.config_path, 'w', encoding='utf-8') as f:
|
|||
|
|
json.dump(default_config, f, ensure_ascii=False, indent=2)
|
|||
|
|
logger.info(f"创建默认配置文件: {self.config_path}")
|
|||
|
|
|
|||
|
|
def save_cookies(self) -> None:
|
|||
|
|
"""保存cookies配置"""
|
|||
|
|
try:
|
|||
|
|
cookies_data = []
|
|||
|
|
for cookie_info in self.cookies.values():
|
|||
|
|
cookies_data.append({
|
|||
|
|
'name': cookie_info.name,
|
|||
|
|
'cookie_string': cookie_info.cookie_string,
|
|||
|
|
'last_used': cookie_info.last_used.isoformat(),
|
|||
|
|
'use_count': cookie_info.use_count,
|
|||
|
|
'is_valid': cookie_info.is_valid,
|
|||
|
|
'failure_count': cookie_info.failure_count,
|
|||
|
|
'user_info': cookie_info.user_info
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
config_data = {'cookies': cookies_data}
|
|||
|
|
os.makedirs(os.path.dirname(self.config_path), exist_ok=True)
|
|||
|
|
with open(self.config_path, 'w', encoding='utf-8') as f:
|
|||
|
|
json.dump(config_data, f, ensure_ascii=False, indent=2)
|
|||
|
|
logger.debug("保存cookies配置成功")
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"保存cookies配置失败: {str(e)}")
|
|||
|
|
|
|||
|
|
def add_cookie(self, name: str, cookie_string: str, user_info: Optional[Dict] = None) -> bool:
|
|||
|
|
"""添加新的cookie"""
|
|||
|
|
if name in self.cookies:
|
|||
|
|
logger.warning(f"Cookie名称 '{name}' 已存在")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
cookie_info = CookieInfo(
|
|||
|
|
name=name,
|
|||
|
|
cookie_string=cookie_string,
|
|||
|
|
user_info=user_info
|
|||
|
|
)
|
|||
|
|
self.cookies[name] = cookie_info
|
|||
|
|
self.save_cookies()
|
|||
|
|
logger.info(f"添加新cookie: {name}")
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
def remove_cookie(self, name: str) -> bool:
|
|||
|
|
"""删除cookie"""
|
|||
|
|
if name not in self.cookies:
|
|||
|
|
logger.warning(f"Cookie名称 '{name}' 不存在")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
del self.cookies[name]
|
|||
|
|
if self.current_cookie == name:
|
|||
|
|
self.current_cookie = None
|
|||
|
|
self.save_cookies()
|
|||
|
|
logger.info(f"删除cookie: {name}")
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
def get_valid_cookies(self) -> List[str]:
|
|||
|
|
"""获取有效的cookie名称列表"""
|
|||
|
|
return [name for name, cookie_info in self.cookies.items()
|
|||
|
|
if cookie_info.is_valid and cookie_info.failure_count < self.max_failures]
|
|||
|
|
|
|||
|
|
def get_cookie(self, name: Optional[str] = None) -> Optional[CookieInfo]:
|
|||
|
|
"""获取指定或最优的cookie"""
|
|||
|
|
if name:
|
|||
|
|
return self.cookies.get(name) if name in self.cookies else None
|
|||
|
|
|
|||
|
|
# 获取最优cookie
|
|||
|
|
valid_cookies = self.get_valid_cookies()
|
|||
|
|
if not valid_cookies:
|
|||
|
|
logger.error("没有可用的cookies")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# 选择策略：优先使用失败次数少、使用次数少的cookie
|
|||
|
|
best_cookie = None
|
|||
|
|
best_score = float('inf')
|
|||
|
|
|
|||
|
|
for cookie_name in valid_cookies:
|
|||
|
|
cookie_info = self.cookies[cookie_name]
|
|||
|
|
# 计算评分：失败次数 * 10 + 使用次数 + 时间因子
|
|||
|
|
time_factor = (datetime.now() - cookie_info.last_used).total_seconds() / 3600 # 小时
|
|||
|
|
score = cookie_info.failure_count * 10 + cookie_info.use_count - time_factor
|
|||
|
|
|
|||
|
|
if score < best_score:
|
|||
|
|
best_score = score
|
|||
|
|
best_cookie = cookie_info
|
|||
|
|
|
|||
|
|
return best_cookie
|
|||
|
|
|
|||
|
|
def get_cookie_string(self, name: Optional[str] = None) -> Optional[str]:
|
|||
|
|
"""获取cookie字符串"""
|
|||
|
|
cookie_info = self.get_cookie(name)
|
|||
|
|
if cookie_info:
|
|||
|
|
# 更新使用信息
|
|||
|
|
cookie_info.last_used = datetime.now()
|
|||
|
|
cookie_info.use_count += 1
|
|||
|
|
self.current_cookie = cookie_info.name
|
|||
|
|
self.save_cookies()
|
|||
|
|
return cookie_info.cookie_string
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def mark_cookie_failure(self, name: Optional[str] = None) -> None:
|
|||
|
|
"""标记cookie使用失败"""
|
|||
|
|
cookie_name = name or self.current_cookie
|
|||
|
|
if cookie_name and cookie_name in self.cookies:
|
|||
|
|
cookie_info = self.cookies[cookie_name]
|
|||
|
|
cookie_info.failure_count += 1
|
|||
|
|
|
|||
|
|
if cookie_info.failure_count >= self.max_failures:
|
|||
|
|
cookie_info.is_valid = False
|
|||
|
|
logger.warning(f"Cookie '{cookie_name}' 已失效")
|
|||
|
|
|
|||
|
|
self.save_cookies()
|
|||
|
|
|
|||
|
|
def mark_cookie_success(self, name: Optional[str] = None) -> None:
|
|||
|
|
"""标记cookie使用成功"""
|
|||
|
|
cookie_name = name or self.current_cookie
|
|||
|
|
if cookie_name and cookie_name in self.cookies:
|
|||
|
|
cookie_info = self.cookies[cookie_name]
|
|||
|
|
cookie_info.failure_count = 0 # 重置失败次数
|
|||
|
|
cookie_info.is_valid = True
|
|||
|
|
self.save_cookies()
|
|||
|
|
|
|||
|
|
def get_random_cookie(self) -> Optional[CookieInfo]:
|
|||
|
|
"""随机获取一个有效的cookie"""
|
|||
|
|
valid_cookies = self.get_valid_cookies()
|
|||
|
|
if not valid_cookies:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
cookie_name = random.choice(valid_cookies)
|
|||
|
|
return self.cookies[cookie_name]
|
|||
|
|
|
|||
|
|
def get_statistics(self) -> CookieStats:
|
|||
|
|
"""获取cookie使用统计"""
|
|||
|
|
total_cookies = len(self.cookies)
|
|||
|
|
valid_cookies = len(self.get_valid_cookies())
|
|||
|
|
invalid_cookies = total_cookies - valid_cookies
|
|||
|
|
|
|||
|
|
cookie_details = []
|
|||
|
|
for name, cookie_info in self.cookies.items():
|
|||
|
|
cookie_details.append({
|
|||
|
|
'name': name,
|
|||
|
|
'is_valid': cookie_info.is_valid,
|
|||
|
|
'use_count': cookie_info.use_count,
|
|||
|
|
'failure_count': cookie_info.failure_count,
|
|||
|
|
'last_used': cookie_info.last_used.isoformat()
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
return CookieStats(
|
|||
|
|
total_cookies=total_cookies,
|
|||
|
|
valid_cookies=valid_cookies,
|
|||
|
|
invalid_cookies=invalid_cookies,
|
|||
|
|
current_cookie=self.current_cookie,
|
|||
|
|
cookie_details=cookie_details
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def refresh_cookie_validity(self) -> None:
|
|||
|
|
"""刷新cookie有效性（重置失败次数）"""
|
|||
|
|
for cookie_info in self.cookies.values():
|
|||
|
|
if cookie_info.failure_count < self.max_failures:
|
|||
|
|
cookie_info.is_valid = True
|
|||
|
|
cookie_info.failure_count = 0
|
|||
|
|
self.save_cookies()
|
|||
|
|
logger.info("已刷新所有cookie的有效性")
|
|||
|
|
|
|||
|
|
def rotate_cookies(self) -> None:
|
|||
|
|
"""轮换cookies（可定期调用）"""
|
|||
|
|
if len(self.get_valid_cookies()) <= 1:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 如果当前cookie使用时间过长，切换到其他cookie
|
|||
|
|
if self.current_cookie:
|
|||
|
|
current_cookie_info = self.cookies.get(self.current_cookie)
|
|||
|
|
if current_cookie_info:
|
|||
|
|
time_since_last_use = (datetime.now() - current_cookie_info.last_used).total_seconds()
|
|||
|
|
if time_since_last_use > self.rotation_interval:
|
|||
|
|
self.current_cookie = None
|
|||
|
|
logger.info("Cookie轮换：切换到其他cookie")
|