MineNasAI/src/minenasai/webtui/auth.py

"""用户认证模块

提供 Token 认证和会话管理
"""

from __future__ import annotations

import hashlib
import hmac
import secrets
import time
from dataclasses import dataclass, field
from typing import Any

from minenasai.core import get_logger, get_settings

logger = get_logger(__name__)


@dataclass
class AuthToken:
    """认证令牌"""

    token: str
    user_id: str
    created_at: float
    expires_at: float
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def is_expired(self) -> bool:
        """是否已过期"""
        return time.time() > self.expires_at

    @property
    def remaining_time(self) -> float:
        """剩余有效时间（秒）"""
        return max(0, self.expires_at - time.time())


class AuthManager:
    """认证管理器"""

    def __init__(self, secret_key: str | None = None) -> None:
        """初始化认证管理器

        Args:
            secret_key: 签名密钥
        """
        self.settings = get_settings()
        self.secret_key = secret_key or self.settings.webtui_secret_key
        self._tokens: dict[str, AuthToken] = {}
        self._user_sessions: dict[str, set[str]] = {}  # user_id -> set of tokens

    def generate_token(
        self,
        user_id: str,
        expires_in: int | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> str:
        """生成访问令牌

        Args:
            user_id: 用户 ID
            expires_in: 有效期（秒），默认使用配置
            metadata: 额外元数据

        Returns:
            访问令牌
        """
        # 生成随机令牌
        random_part = secrets.token_urlsafe(24)
        timestamp = str(int(time.time()))

        # 签名
        signature = self._sign(f"{random_part}:{timestamp}:{user_id}")
        token = f"{random_part}.{timestamp}.{signature[:16]}"

        # 有效期
        if expires_in is None:
            expires_in = self.settings.webtui.session_timeout

        now = time.time()
        auth_token = AuthToken(
            token=token,
            user_id=user_id,
            created_at=now,
            expires_at=now + expires_in,
            metadata=metadata or {},
        )

        # 存储
        self._tokens[token] = auth_token

        # 记录用户会话
        if user_id not in self._user_sessions:
            self._user_sessions[user_id] = set()
        self._user_sessions[user_id].add(token)

        logger.info("生成令牌", user_id=user_id, expires_in=expires_in)
        return token

    def verify_token(self, token: str) -> AuthToken | None:
        """验证令牌

        Args:
            token: 访问令牌

        Returns:
            AuthToken 或 None（无效时）
        """
        auth_token = self._tokens.get(token)

        if auth_token is None:
            logger.debug("令牌不存在", token=token[:20] + "...")
            return None

        if auth_token.is_expired:
            logger.debug("令牌已过期", token=token[:20] + "...")
            self.revoke_token(token)
            return None

        return auth_token

    def revoke_token(self, token: str) -> bool:
        """撤销令牌

        Args:
            token: 访问令牌

        Returns:
            是否成功撤销
        """
        auth_token = self._tokens.pop(token, None)
        if auth_token:
            # 从用户会话中移除
            if auth_token.user_id in self._user_sessions:
                self._user_sessions[auth_token.user_id].discard(token)
            logger.info("撤销令牌", user_id=auth_token.user_id)
            return True
        return False

    def revoke_user_tokens(self, user_id: str) -> int:
        """撤销用户的所有令牌

        Args:
            user_id: 用户 ID

        Returns:
            撤销的令牌数量
        """
        tokens = self._user_sessions.pop(user_id, set())
        for token in tokens:
            self._tokens.pop(token, None)

        if tokens:
            logger.info("撤销用户所有令牌", user_id=user_id, count=len(tokens))

        return len(tokens)

    def refresh_token(self, token: str, extends_by: int | None = None) -> str | None:
        """刷新令牌

        Args:
            token: 当前令牌
            extends_by: 延长时间（秒）

        Returns:
            新令牌或 None（失败时）
        """
        auth_token = self.verify_token(token)
        if auth_token is None:
            return None

        # 生成新令牌
        new_token = self.generate_token(
            user_id=auth_token.user_id,
            expires_in=extends_by,
            metadata=auth_token.metadata,
        )

        # 撤销旧令牌
        self.revoke_token(token)

        return new_token

    def cleanup_expired(self) -> int:
        """清理过期令牌

        Returns:
            清理的令牌数量
        """
        now = time.time()
        expired = [
            token for token, auth_token in self._tokens.items()
            if auth_token.expires_at < now
        ]

        for token in expired:
            self.revoke_token(token)

        if expired:
            logger.info("清理过期令牌", count=len(expired))

        return len(expired)

    def _sign(self, data: str) -> str:
        """生成签名"""
        return hmac.new(
            self.secret_key.encode(),
            data.encode(),
            hashlib.sha256
        ).hexdigest()

    def get_stats(self) -> dict[str, Any]:
        """获取统计信息"""
        return {
            "total_tokens": len(self._tokens),
            "total_users": len(self._user_sessions),
            "tokens_per_user": {
                user_id: len(tokens)
                for user_id, tokens in self._user_sessions.items()
            },
        }


# 全局认证管理器
_auth_manager: AuthManager | None = None


def get_auth_manager() -> AuthManager:
    """获取全局认证管理器"""
    global _auth_manager
    if _auth_manager is None:
        _auth_manager = AuthManager()
    return _auth_manager