test/app/core/security.py

# -*- coding: utf-8 -*-
"""
安全相关功能
包含鉴权、token验证等安全功能
"""

from typing import Optional
from fastapi import Request
from .config import settings

TOKEN_HEADER_NAME = "X-NLS-Token"
AUTH_OPTIONAL_PLACEHOLDER = "optional"
WEBSOCKET_QUERY_TOKEN_KEYS = ("token", "x_nls_token", "X-NLS-Token")


def normalize_token(token: Optional[str]) -> Optional[str]:
    """将 token 归一化为非空字符串或 None。"""
    if token is None:
        return None

    normalized = token.strip()
    return normalized or None


def get_expected_api_key(expected_token: Optional[str] = None) -> Optional[str]:
    """获取归一化后的期望 API_KEY。"""
    if expected_token is not None:
        return normalize_token(expected_token)
    return normalize_token(settings.API_KEY)


def mask_sensitive_data(
    data: str, mask_char: str = "*", keep_prefix: int = 4, keep_suffix: int = 4
) -> str:
    """遮盖敏感数据

    Args:
        data: 需要遮盖的数据
        mask_char: 遮盖字符
        keep_prefix: 保留前缀字符数
        keep_suffix: 保留后缀字符数

    Returns:
        遮盖后的数据
    """
    if not data or len(data) <= keep_prefix + keep_suffix:
        return data

    prefix = data[:keep_prefix]
    suffix = data[-keep_suffix:] if keep_suffix > 0 else ""
    mask_length = len(data) - keep_prefix - keep_suffix
    mask = mask_char * mask_length

    return f"{prefix}{mask}{suffix}"


def validate_token_value(token: Optional[str], expected_token: Optional[str] = None) -> bool:
    """验证访问令牌

    Args:
        token: 客户端提供的token
        expected_token: 期望的token值（从环境变量读取），如果为None则鉴权可选

    Returns:
        bool: 验证结果
    """
    normalized_expected_token = get_expected_api_key(expected_token)
    if not normalized_expected_token:
        return True

    normalized_token = normalize_token(token)
    if not normalized_token:
        return False

    # 简单的token格式验证（长度检查）
    if len(normalized_token) < 10:
        return False

    # 验证token是否匹配
    if normalized_token != normalized_expected_token:
        return False

    return True


def extract_header_token(request: Request) -> Optional[str]:
    """从标准头部提取 token。"""
    return normalize_token(request.headers.get(TOKEN_HEADER_NAME))


def extract_bearer_token(request: Request) -> Optional[str]:
    """从 Authorization: Bearer 提取 token。"""
    auth_header = request.headers.get("Authorization")
    if not auth_header:
        return None

    scheme, _, value = auth_header.partition(" ")
    if scheme.lower() != "bearer":
        return None
    return normalize_token(value)


def extract_openai_token(request: Request) -> Optional[str]:
    """OpenAI 兼容接口鉴权：优先 Bearer，其次 X-NLS-Token。"""
    return extract_bearer_token(request) or extract_header_token(request)


def extract_websocket_token(websocket) -> Optional[str]:
    """从 WebSocket 连接中提取 token。"""
    if hasattr(websocket, "headers"):
        token = normalize_token(websocket.headers.get(TOKEN_HEADER_NAME))
        if token:
            return token

    if hasattr(websocket, "query_params"):
        for key in WEBSOCKET_QUERY_TOKEN_KEYS:
            token = normalize_token(websocket.query_params.get(key))
            if token:
                return token

    return None


def _validate_resolved_token(
    token: Optional[str],
    missing_message: str,
    expected_token: Optional[str] = None,
) -> tuple[bool, str]:
    """统一 token 校验逻辑。"""
    expected = get_expected_api_key(expected_token)
    normalized_token = normalize_token(token)

    if not expected:
        return True, normalized_token or AUTH_OPTIONAL_PLACEHOLDER

    if not normalized_token:
        return False, missing_message

    if not validate_token_value(normalized_token, expected):
        masked_token = mask_sensitive_data(normalized_token)
        return False, f"Gateway:ACCESS_DENIED:The token '{masked_token}' is invalid!"

    return True, normalized_token


def validate_token(request: Request, task_id: str = "") -> tuple[bool, str]:
    """验证X-NLS-Token头部"""
    _ = task_id
    token = extract_header_token(request)
    return _validate_resolved_token(token, "缺少X-NLS-Token头部")


def validate_openai_token(request: Request, task_id: str = "") -> tuple[bool, str]:
    """验证 OpenAI 兼容接口 token（Bearer/X-NLS-Token）。"""
    _ = task_id
    token = extract_openai_token(request)
    return _validate_resolved_token(token, "缺少Authorization Bearer或X-NLS-Token头部")


def validate_token_websocket(token: str, task_id: str = "") -> tuple[bool, str]:
    """验证WebSocket连接中的token"""
    _ = task_id
    return _validate_resolved_token(token, "缺少token参数")


def validate_websocket_token(websocket, task_id: str = "") -> tuple[bool, str]:
    """验证 WebSocket 连接 token（header/query 参数）。"""
    _ = task_id
    token = extract_websocket_token(websocket)
    return _validate_resolved_token(
        token,
        "缺少鉴权信息，请通过 X-NLS-Token header 或 token/x_nls_token 查询参数传入",
    )