shudao-main/shudao-chat-py/utils/token.py

"""
本地 Token 验证工具
用于区分本地生成的 token 和外部系统的 token
"""
import jwt
from typing import Optional
from utils.logger import logger


def verify_local_token(token: str) -> Optional[dict]:
    """
    验证是否为本地生成的 token

    Args:
        token: JWT token 字符串

    Returns:
        如果是本地 token 返回解码后的数据，否则返回 None
    """
    if not token:
        return None

    try:
        # 尝试解码 token（不验证签名，只检查格式）
        # 本地 token 应该包含特定的字段，如 account, username 等
        decoded = jwt.decode(token, options={"verify_signature": False})

        # 将解码的 token 打印出来以供调试分析
        logger.info(f"[Token验证] 解码后的 Token 负载: {decoded}")

        # 检查是否包含本地 token 的特征字段
        # 或者包含 user_id, id, sub, sub_id, name 等 (兼容各种其他系统的 token 格式)
        if any(k in decoded for k in ["account", "username", "user_id", "id", "sub", "userId", "name", "email", "uid"]):
            # 尽可能提取出唯一的用户名/标识
            username = (
                decoded.get('username') or
                decoded.get('account') or
                decoded.get('name') or
                decoded.get('email') or
                f"User_{decoded.get('user_id', decoded.get('id', decoded.get('sub', decoded.get('uid', 'unknown'))))}"
            )

            # 补全缺失的关键字段，避免后续代码报错
            if 'username' not in decoded:
                decoded['username'] = username
            if 'account' not in decoded:
                decoded['account'] = username
            
            # 统一 user_id 字段（兼容 userCode）
            if 'userCode' in decoded and 'user_id' not in decoded:
                decoded['user_id'] = decoded['userCode']
            elif 'user_id' in decoded and 'userCode' not in decoded:
                decoded['userCode'] = decoded['user_id']
            elif 'id' in decoded:
                if 'user_id' not in decoded:
                    decoded['user_id'] = decoded['id']
                if 'userCode' not in decoded:
                    decoded['userCode'] = decoded['id']
            elif 'sub' in decoded:
                if 'user_id' not in decoded:
                    decoded['user_id'] = decoded['sub']
                if 'userCode' not in decoded:
                    decoded['userCode'] = decoded['sub']

            logger.info(f"[Token验证] 识别为有效 token: {username}")
            return decoded

        # 如果以上所有字段都没有，但它是个合法的字典结构，我们也强行给它通过（作为游客）
        if isinstance(decoded, dict):
            logger.info("[Token验证] 未找到明确用户字段，作为匿名用户处理")
            decoded['username'] = "Anonymous"
            decoded['account'] = "Anonymous"
            decoded['id'] = 0
            return decoded

        logger.info("[Token验证] 不是本地 token 格式")
        return None

    except jwt.DecodeError:
        logger.info("[Token验证] Token 解码失败，不是有效的 JWT")
        return None
    except Exception as e:
        logger.warning(f"[Token验证] Token 验证异常: {e}")
        return None


def is_local_token(token: str) -> bool:
    """
    判断是否为本地 token

    Args:
        token: JWT token 字符串

    Returns:
        True 表示本地 token，False 表示外部 token
    """
    return verify_local_token(token) is not None