status_manager.py

import asyncio
import json
import os
from typing import Dict, List, Set

class StatusManager:
    def __init__(self, storage_file: str = "cookie_status.json"):
        self.storage_file = storage_file
        self.total_requests = 0
        self.success_counts = 0
        self.failure_counts = 0
        self.expired_cookies: Set[str] = set()
        self.lock = asyncio.Lock()
        self._load_status()

    def _load_status(self):
        """从 JSON 文件加载失效 Cookie"""
        if os.path.exists(self.storage_file):
            try:
                with open(self.storage_file, "r", encoding="utf-8") as f:
                    data = json.load(f)
                    self.expired_cookies = set(data.get("expired_cookies", []))
                    self.total_requests = data.get("total_requests", 0)
                    self.success_counts = data.get("success_counts", 0)
                    self.failure_counts = data.get("failure_counts", 0)
            except Exception:
                pass

    def _save_status(self):
        """保存状态到 JSON 文件"""
        try:
            with open(self.storage_file, "w", encoding="utf-8") as f:
                json.dump({
                    "expired_cookies": list(self.expired_cookies),
                    "total_requests": self.total_requests,
                    "success_counts": self.success_counts,
                    "failure_counts": self.failure_counts
                }, f, ensure_ascii=False, indent=2)
        except Exception:
            pass

    async def mark_cookie_expired(self, cookie: str):
        """标记 Cookie 为失效"""
        async with self.lock:
            self.expired_cookies.add(cookie)
            self._save_status()

    def is_cookie_expired(self, cookie: str) -> bool:
        """检查 Cookie 是否失效"""
        return cookie in self.expired_cookies

    async def record_request(self, success: bool = True):
        """记录请求"""
        async with self.lock:
            self.total_requests += 1
            if success:
                self.success_counts += 1
            else:
                self.failure_counts += 1
            # 频繁写入可能影响性能，但在低并发下影响不大
            self._save_status()

    async def get_stats(self) -> Dict:
        """获取统计信息"""
        async with self.lock:
            return {
                "total_requests": self.total_requests,
                "success_counts": self.success_counts,
                "failure_counts": self.failure_counts,
                "expired_count": len(self.expired_cookies)
            }

    async def get_expired_cookies(self) -> List[str]:
        """获取所有已失效的 Cookie"""
        async with self.lock:
            return list(self.expired_cookies)

    async def reset_expired_cookies(self):
        """重置所有失效 Cookie"""
        async with self.lock:
            self.expired_cookies.clear()
            self._save_status()

# 全局单例
status_manager = StatusManager()