auth_manager.py

import asyncio
import json
import logging
import os
from typing import Dict, List, Tuple

import httpx

import config
import fingerprint
import timing

logger = logging.getLogger(__name__)

def _mask_email(email: str) -> str:
    parts = email.split("@")
    if len(parts) == 2:
        return f"{parts[0][:2]}***@{parts[1]}"
    return "***"

class AuthManager:
    def __init__(self, status_file: str = "accounts_status.json"):
        self.status_file = status_file
        self.lock = asyncio.Lock()
        
        # Parse accounts
        self.accounts: List[Tuple[str, str]] = []  # [(email, password), ...]
        for acc in config.ONYX_ACCOUNTS:
            parts = acc.split(":", 1)
            if len(parts) == 2:
                self.accounts.append((parts[0].strip(), parts[1].strip()))
                
        # Load saved status
        # status: { email: { "fastapiusersauth": "...", "fastapiusersoauthcsrf": "...", "valid": True } }
        self.status: Dict[str, Dict] = {}
        self._load_status()
        
        # For legacy cookie support
        self.legacy_cookies = config.ONYX_AUTH_COOKIES.copy()
        
        # Round robin index
        self._account_index = 0
        self._legacy_index = 0

    def _load_status(self):
        if os.path.exists(self.status_file):
            try:
                with open(self.status_file, "r", encoding="utf-8") as f:
                    self.status = json.load(f)
            except Exception as e:
                logger.warning(f"Failed to load accounts status: {e}")

    def _save_status(self):
        try:
            with open(self.status_file, "w", encoding="utf-8") as f:
                json.dump(self.status, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.error(f"Failed to save accounts status: {e}")

    async def get_valid_cookie(self, client: httpx.AsyncClient) -> Tuple[str, str, str]:
        """返回有效的 (fastapiusersauth, fastapiusersoauthcsrf, email) 三元组。
        
        整个选择+验证过程在 asyncio.Lock 内完成，避免并发请求同时选中
        同一个账号并发登录导致 token 互相覆盖。asyncio.Lock 在 await 期间
        会释放事件循环控制权，不会像 threading.Lock 那样阻塞整个进程。
        """
        async with self.lock:
            # If we don't have accounts but have legacy cookies, fallback
            if not self.accounts:
                if not self.legacy_cookies:
                    raise RuntimeError("No accounts and no legacy cookies configured. Please check config.")
                c = self.legacy_cookies[self._legacy_index % len(self.legacy_cookies)]
                self._legacy_index += 1
                return c, "", ""  # 旧版 cookie 没有 CSRF token 及 email

            for _ in range(len(self.accounts)):
                email, password = self.accounts[self._account_index % len(self.accounts)]
                # Move to next account for the *next* request
                self._account_index += 1
                acc_status = self.status.get(email, {})
                
                # Skip accounts that are permanently banned or quota exhausted
                if acc_status.get("reason") == "额度耗尽 / 被拒绝":
                    continue
                    
                # If the account currently has a valid cookie, use it
                if acc_status.get("valid") and acc_status.get("fastapiusersauth"):
                    return acc_status["fastapiusersauth"], acc_status.get("fastapiusersoauthcsrf", ""), email
                    
                # Account has no cookie or its cookie is invalid — try login/refresh
                success = await self._ensure_account_valid(client, email, password)
                if success:
                    return self.status[email]["fastapiusersauth"], self.status[email].get("fastapiusersoauthcsrf", ""), email
                
                # Login/refresh failed, move on to the next account
                    
            raise RuntimeError("Failed to obtain any valid authentication cookie from configured accounts.")

    async def _ensure_account_valid(self, client: httpx.AsyncClient, email: str, password: str) -> bool:
        """Ensures the account has valid cookies, trying refresh or login if needed.
        NOTE: Called from within self.lock, no need for additional locking."""
        acc_status = self.status.get(email, {})
        auth = acc_status.get("fastapiusersauth")
        csrf = acc_status.get("fastapiusersoauthcsrf")
            
        if auth and csrf:
            # Try refresh
            success = await self._refresh(client, email, auth, csrf)
            if success:
                return True
                
        # If refresh fails or no tokens, do full login
        return await self._login(client, email, password)

    async def _login(self, client: httpx.AsyncClient, email: str, password: str) -> bool:
        logger.info(f"Attempting full login for account {_mask_email(email)}")
        persona = fingerprint.get_persona_for_account(email)
        headers = fingerprint.get_base_headers(persona)
        headers.update({
            "Origin": config.ONYX_BASE_URL,
            "Referer": f"{config.ONYX_BASE_URL}/login"
        })
        
        try:
            # Step 1: Get CSRF token
            auth_url = f"{config.ONYX_BASE_URL}/api/auth/oauth/authorize"
            r1 = await client.get(auth_url, headers=headers, follow_redirects=False)
            csrf = r1.cookies.get("fastapiusersoauthcsrf")
            
            await timing.micro_delay("click")
            
            # Step 2: Login
            login_url = f"{config.ONYX_BASE_URL}/api/auth/login"
            data = {"username": email, "password": password}
            headers["Content-Type"] = "application/x-www-form-urlencoded"
            
            cookies = {}
            if csrf:
                cookies["fastapiusersoauthcsrf"] = csrf
                
            r2 = await client.post(login_url, data=data, headers=headers, cookies=cookies)
            
            if r2.status_code in [200, 204]:
                new_auth = r2.cookies.get("fastapiusersauth")
                new_csrf = r2.cookies.get("fastapiusersoauthcsrf") or csrf
                
                if new_auth:
                    self.status[email] = {
                        "fastapiusersauth": new_auth,
                        "fastapiusersoauthcsrf": new_csrf,
                        "valid": True
                    }
                    self._save_status()
                    logger.info(f"Successfully logged in {_mask_email(email)}")
                    return True
            
            logger.error(f"Login failed for {_mask_email(email)}: HTTP {r2.status_code}")
            return False
            
        except Exception as e:
            logger.error(f"Error during login for {_mask_email(email)}: {e}")
            return False

    async def _refresh(self, client: httpx.AsyncClient, email: str, auth: str, csrf: str) -> bool:
        logger.info(f"Attempting token refresh for account {_mask_email(email)}")
        persona = fingerprint.get_persona_for_account(email)
        headers = fingerprint.get_base_headers(persona)
        headers.update({
            "Origin": config.ONYX_BASE_URL,
            "Referer": f"{config.ONYX_BASE_URL}/login"
        })
        cookies = {
            "fastapiusersauth": auth,
            "fastapiusersoauthcsrf": csrf
        }
        refresh_url = f"{config.ONYX_BASE_URL}/api/auth/refresh"
        
        try:
            r = await client.post(refresh_url, headers=headers, cookies=cookies)
            if r.status_code in [200, 204]:
                new_auth = r.cookies.get("fastapiusersauth")
                new_csrf = r.cookies.get("fastapiusersoauthcsrf") or csrf
                if new_auth:
                    self.status[email] = {
                        "fastapiusersauth": new_auth,
                        "fastapiusersoauthcsrf": new_csrf,
                        "valid": True
                    }
                    self._save_status()
                    logger.info(f"Successfully refreshed token for {_mask_email(email)}")
                    return True
                    
            logger.warning(f"Refresh failed for {_mask_email(email)} (HTTP {r.status_code}), will attempt full login.")
            return False
        except Exception as e:
            logger.error(f"Error during refresh for {_mask_email(email)}: {e}")
            return False

    async def report_unauthorized(self, client: httpx.AsyncClient, auth_cookie: str):
        """Called when an API request returns 401/403 with a specific cookie.
        Marks it as invalid and tries to refresh it or login."""
        async with self.lock:
            if auth_cookie in self.legacy_cookies:
                self.legacy_cookies.remove(auth_cookie)
                logger.warning("Removed legacy cookie due to 401 Unauthorized")
                return

            target_email = None
            target_password = None
            target_csrf = None

            for email, stats in self.status.items():
                auth_val = self._extract_auth_value(auth_cookie)
                stored_val = self._extract_auth_value(stats.get("fastapiusersauth", ""))
                
                if stored_val and stored_val == auth_val:
                    stats["valid"] = False
                    self._save_status()
                    target_email = email
                    target_csrf = stats.get("fastapiusersoauthcsrf")
                    break
                    
            if target_email:
                for email, pwd in self.accounts:
                    if email == target_email:
                        target_password = pwd
                        break

            if target_email and target_password:
                success = False
                if target_csrf:
                    success = await self._refresh(client, target_email, auth_cookie, target_csrf)
                if not success:
                    await self._login(client, target_email, target_password)

    async def report_forbidden(self, client: httpx.AsyncClient, auth_cookie: str):
        """收到 403 Forbidden 时调用。
        先尝试重新登录获取新 cookie，如果重新登录成功则恢复账号，
        只有确实无法恢复时才标记为额度耗尽。"""
        async with self.lock:
            if auth_cookie in self.legacy_cookies:
                self.legacy_cookies.remove(auth_cookie)
                logger.warning("移除了一个旧版 cookie (403 Forbidden)")
                return

            target_email = None
            target_password = None

            for email, stats in self.status.items():
                auth_val = self._extract_auth_value(auth_cookie)
                stored_val = self._extract_auth_value(stats.get("fastapiusersauth", ""))
                
                if stored_val and stored_val == auth_val:
                    # 先标记为无效，但不写 reason（等重新登录结果决定）
                    stats["valid"] = False
                    self._save_status()
                    target_email = email
                    break
                    
            if target_email:
                for email, pwd in self.accounts:
                    if email == target_email:
                        target_password = pwd
                        break

            # 尝试重新登录恢复账号
            if target_email and target_password:
                logger.info(f"403 后尝试重新登录账号 {_mask_email(target_email)}...")
                success = await self._login(client, target_email, target_password)
                if success:
                    logger.info(f"账号 {_mask_email(target_email)} 重新登录成功，可能只是 cookie 过期而非额度耗尽")
                    return  # 账号已恢复，_login 已设置 valid=True
                else:
                    logger.warning(f"账号 {_mask_email(target_email)} 重新登录也失败了，标记为额度耗尽")
            
            # 如果重新登录失败或找不到对应账号，标记为额度耗尽
            if target_email and target_email in self.status:
                self.status[target_email]["valid"] = False
                self.status[target_email]["reason"] = "额度耗尽 / 被拒绝"
                self._save_status()

    def _extract_auth_value(self, cookie_str: str) -> str:
        if not cookie_str:
            return ""
        if "fastapiusersauth=" in cookie_str:
            for piece in cookie_str.split(";"):
                piece = piece.strip()
                if piece.startswith("fastapiusersauth="):
                    return piece.split("=", 1)[1]
        return cookie_str.strip()

auth_manager = AuthManager()