assistant/agent/utils/token_manager.py

"""
Token Manager - 管理 JWT token 的获取和使用

支持从 Chatwoot contact custom_attributes 中获取用户的 JWT token
"""
from typing import Optional
from utils.logger import get_logger

logger = get_logger(__name__)


class TokenManager:
    """管理用户 JWT token"""

    @staticmethod
    def extract_token_from_sender(sender: Optional[dict]) -> Optional[str]:
        """从 sender 对象中提取 JWT token

        支持从以下位置提取 token（按优先级排序）：
        1. sender.jwt_token（根级别）
        2. sender.mall_token（根级别）
        3. sender.custom_attributes.jwt_token
        4. sender.custom_attributes.mall_token
        5. sender.custom_attributes.access_token
        6. sender.custom_attributes.auth_token
        7. sender.custom_attributes.token

        Args:
            sender: Chatwoot sender 对象（来自 conversation.meta.sender）

        Returns:
            JWT token 字符串，如果未找到则返回 None
        """
        if not sender:
            logger.debug("No sender provided")
            return None

        # 1. 优先从根级别获取 token
        root_token = sender.get("jwt_token") or sender.get("mall_token")
        if root_token:
            logger.debug("JWT token found at sender root level")
            logger.debug(f"Token prefix: {root_token[:20]}...")
            return root_token

        # 2. 从 custom_attributes 中获取 token
        return TokenManager.extract_token_from_contact(sender)

    @staticmethod
    def extract_token_from_contact(contact: Optional[dict]) -> Optional[str]:
        """从 Chatwoot contact 中提取 JWT token

        Args:
            contact: Chatwoot contact 对象，包含 custom_attributes

        Returns:
            JWT token 字符串，如果未找到则返回 None
        """
        if not contact:
            logger.debug("No contact provided")
            return None

        # 从 custom_attributes 中获取 token
        custom_attributes = contact.get("custom_attributes", {})
        if not custom_attributes:
            logger.debug("No custom_attributes in contact")
            return None

        # 尝试多种可能的字段名
        token = (
            custom_attributes.get("jwt_token") or
            custom_attributes.get("mall_token") or
            custom_attributes.get("access_token") or
            custom_attributes.get("auth_token") or
            custom_attributes.get("token")
        )

        if token:
            logger.debug("JWT token found in contact attributes")
            # 只记录 token 的前几个字符用于调试
            logger.debug(f"Token prefix: {token[:20]}...")
        else:
            logger.debug("No JWT token found in contact custom_attributes")

        return token

    @staticmethod
    def validate_token(token: str) -> bool:
        """验证 token 格式是否有效

        Args:
            token: JWT token 字符串

        Returns:
            True 如果 token 格式有效
        """
        if not token or not isinstance(token, str):
            return False

        # JWT token 通常是 header.payload.signature 格式
        parts = token.split(".")
        if len(parts) != 3:
            logger.warning("Invalid JWT token format")
            return False

        return True

    @staticmethod
    def get_token_from_context(context: dict, contact: Optional[dict] = None) -> Optional[str]:
        """从上下文或 contact 中获取 token

        优先级：context > contact

        Args:
            context: 对话上下文
            contact: Chatwoot contact 对象

        Returns:
            JWT token 或 None
        """
        # 首先尝试从 context 中获取（可能之前的对话中已经获取）
        token = context.get("user_token")
        if token and TokenManager.validate_token(token):
            logger.debug("Using token from context")
            return token

        # 其次尝试从 contact 中获取
        if contact:
            token = TokenManager.extract_token_from_contact(contact)
            if token and TokenManager.validate_token(token):
                logger.debug("Using token from contact")
                return token

        logger.debug("No valid JWT token found")
        return None


# 全局 token 管理器
token_manager = TokenManager()