ForcePilot/backend/package/yuxi/utils/crypto.py

"""敏感字段加密与脱敏通用工具。

基于 Fernet 对称加密，支持递归处理配置字典中的敏感字段。
所有加密、解密、脱敏、密钥校验逻辑均封装在 ``CryptoHelper`` 类中，
模块级便捷函数仅作为默认实例的薄封装，便于全局调用与测试注入。
"""

from __future__ import annotations

import os
import re
from typing import Any, Literal

from cryptography.fernet import Fernet, InvalidToken

from yuxi.utils import logger


class CryptoHelper:
    """敏感字段加密与脱敏工具（完整闭环）。

    基于 Fernet 对称加密，支持递归处理配置字典中的敏感字段。

    密钥来源优先级：
        1. 构造函数注入的 ``fernet`` 实例
        2. 构造函数注入的 ``key`` 字符串
        3. 环境变量 ``ENCRYPTION_KEY``

    未提供密钥时，加密/解密操作将透传原值（开发环境兼容），并在首次
    访问时记录一次警告日志；生产环境可通过 ``validate_for_environment``
    强制校验密钥存在性。

    敏感字段判定：字段名（小写归一化后）命中 ``_SENSITIVE_KEYWORDS`` 中
    任一关键词即视为敏感。设计原则是宁可漏判不可误判，因此避免使用过于
    宽泛的词（如 "key" / "cert" / "auth" / "id"），防止误伤
    ``auth_type`` / ``auth_config`` / ``content-type`` / ``client_id``
    等非敏感字段。
    """

    # 加密密钥环境变量名
    _ENV_KEY_NAME = "ENCRYPTION_KEY"

    # 用于识别 Fernet 加密值的统一前缀
    _ENCRYPTION_PREFIX = "enc:"

    # 敏感字段关键词：字段名（小写归一化后）命中任一关键词即视为敏感。
    # 设计原则：宁可漏判（漏判可由调用方补充），不可误判（误判会破坏正常字段）。
    # 因此避免使用过于宽泛的词（如 "key" / "cert" / "auth" / "id"），
    # 防止误伤 auth_type / auth_config / content-type / client_id 等非敏感字段。
    _SENSITIVE_KEYWORDS = (
        # 密码类
        "password",
        "passwd",
        "pwd",
        # 令牌类
        "token",
        "access_token",
        "refresh_token",
        "id_token",
        "bearer",
        "authorization",
        # 密钥类
        "secret",
        "api_key",
        "apikey",
        "api_secret",
        "access_key",
        "secret_key",
        "private_key",
        "client_secret",
        "client_key",
        # 证书类
        "ca_cert",
        "client_cert",
        "certificate",
        # 凭证类
        "credential",
        "credentials",
    )

    def __init__(self, key: str | None = None, fernet: Fernet | None = None) -> None:
        """初始化加密助手。

        Args:
            key: Fernet 密钥字符串。提供时优先使用，忽略环境变量。
            fernet: 已构造的 Fernet 实例。提供时优先于 ``key``。
        """
        if fernet is not None:
            self._fernet: Fernet | None = fernet
            self._injected_key: str | None = None
        elif key is not None:
            self._fernet = Fernet(key.encode())
            self._injected_key = key
        else:
            self._fernet = None
            self._injected_key = None
        self._missing_logged = False

    # ------------------------------------------------------------------
    # 密钥管理
    # ------------------------------------------------------------------

    @classmethod
    def from_env(cls) -> CryptoHelper:
        """从环境变量 ``ENCRYPTION_KEY`` 创建实例。

        环境变量未设置时返回未配置密钥的实例（开发环境兼容）。
        """
        key = os.getenv(cls._ENV_KEY_NAME)
        return cls(key=key) if key else cls()

    def has_key(self) -> bool:
        """检查是否存在有效密钥（不触发日志，不抛异常）。"""
        if self._fernet is not None:
            return True
        if self._injected_key is not None:
            return True
        return bool(os.getenv(self._ENV_KEY_NAME))

    def _get_fernet(self) -> Fernet | None:
        """获取 Fernet 实例，按密钥来源优先级解析。

        密钥未配置时记录一次警告并返回 None；密钥格式无效时抛出 ValueError。
        """
        if self._fernet is not None:
            return self._fernet

        key = self._injected_key if self._injected_key is not None else os.getenv(self._ENV_KEY_NAME)
        if not key:
            if not self._missing_logged:
                logger.warning(
                    f"环境变量 {self._ENV_KEY_NAME} 未设置，配置中的敏感字段将以明文存储。"
                    "生产环境请使用 Fernet.generate_key() 生成密钥并配置。"
                )
                self._missing_logged = True
            return None

        try:
            self._fernet = Fernet(key.encode())
        except Exception as exc:
            raise ValueError(f"{self._ENV_KEY_NAME} 格式无效，无法初始化 Fernet") from exc
        return self._fernet

    def validate_for_environment(self, environment: str) -> None:
        """校验生产环境是否正确配置加密密钥。

        非生产环境跳过校验；生产环境要求密钥必须存在且能初始化 Fernet，
        否则抛出 ValueError。
        """
        if environment != "production":
            return

        if not self.has_key():
            raise ValueError(f"生产环境必须设置 {self._ENV_KEY_NAME}")

        try:
            self._get_fernet()
        except ValueError as exc:
            raise ValueError(f"{self._ENV_KEY_NAME} 格式无效") from exc

    # ------------------------------------------------------------------
    # 单值加密/解密
    # ------------------------------------------------------------------

    def encrypt(self, value: str) -> str:
        """加密字符串，返回带 ``enc:`` 前缀的密文。

        密钥未配置时透传原值。
        """
        fernet = self._get_fernet()
        if fernet is None:
            return value
        encrypted = fernet.encrypt(value.encode())
        return f"{self._ENCRYPTION_PREFIX}{encrypted.decode()}"

    def decrypt(self, value: str) -> str:
        """解密带 ``enc:`` 前缀的密文，返回明文。

        非 ``enc:`` 前缀的值视为未加密，透传返回；密钥未配置时记录警告
        并透传；解密失败（密钥变更或数据损坏）抛出 ValueError。
        """
        if not isinstance(value, str) or not value.startswith(self._ENCRYPTION_PREFIX):
            return value

        fernet = self._get_fernet()
        if fernet is None:
            logger.warning("缺少加密密钥，无法解密配置中的敏感字段")
            return value

        try:
            decrypted = fernet.decrypt(value[len(self._ENCRYPTION_PREFIX) :].encode())
        except InvalidToken as exc:
            raise ValueError("配置中的敏感字段解密失败，密钥可能已变更") from exc
        return decrypted.decode()

    # ------------------------------------------------------------------
    # 递归加密/解密/脱敏
    # ------------------------------------------------------------------

    def encrypt_sensitive_fields(self, config: dict[str, Any] | None) -> dict[str, Any] | None:
        """递归加密配置字典中的敏感字符串字段。仅加密敏感 key 对应的字符串。"""
        return self._transform_config(config, "encrypt")

    def decrypt_sensitive_fields(self, config: dict[str, Any] | None) -> dict[str, Any] | None:
        """递归解密配置字典中的敏感字符串字段。"""
        return self._transform_config(config, "decrypt")

    def mask_sensitive_fields(self, config: dict[str, Any] | None) -> dict[str, Any]:
        """返回配置字典的脱敏副本，用于 API 响应。"""
        result = self._transform_config(config, "mask")
        return result if result is not None else {}

    def mask_all_string_values(self, obj: Any) -> Any:
        """递归将对象中的所有字符串值替换为脱敏占位符。"""
        if isinstance(obj, dict):
            return {k: self.mask_all_string_values(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [self.mask_all_string_values(item) for item in obj]
        if isinstance(obj, str):
            return "***"
        return obj

    # ------------------------------------------------------------------
    # 内部实现
    # ------------------------------------------------------------------

    @classmethod
    def _is_sensitive_key(cls, key: str) -> bool:
        """判断字段名是否属于敏感字段。"""
        normalized = re.sub(r"[^a-z0-9_]+", "_", key.lower())
        return any(keyword in normalized for keyword in cls._SENSITIVE_KEYWORDS)

    def _transform_config(
        self,
        config: dict[str, Any] | None,
        mode: Literal["encrypt", "decrypt", "mask"],
    ) -> dict[str, Any] | None:
        """通用配置转换：递归加密/解密/脱敏敏感字段。"""
        if config is None:
            return None

        if mode == "encrypt":
            return self._encrypt_nested(config)
        if mode == "decrypt":
            return self._decrypt_nested(config)
        # mask
        return self._mask_nested(config)

    def _encrypt_nested(self, obj: Any) -> Any:
        if isinstance(obj, dict):
            return {
                k: self.encrypt(v) if self._is_sensitive_key(k) and isinstance(v, str) else self._encrypt_nested(v)
                for k, v in obj.items()
            }
        if isinstance(obj, list):
            return [self._encrypt_nested(item) for item in obj]
        return obj

    def _decrypt_nested(self, obj: Any) -> Any:
        if isinstance(obj, dict):
            return {
                k: self.decrypt(v) if self._is_sensitive_key(k) and isinstance(v, str) else self._decrypt_nested(v)
                for k, v in obj.items()
            }
        if isinstance(obj, list):
            return [self._decrypt_nested(item) for item in obj]
        return obj

    def _mask_nested(self, obj: Any) -> Any:
        if isinstance(obj, dict):
            return {
                k: self._mask_string(v) if self._is_sensitive_key(k) and isinstance(v, str) else self._mask_nested(v)
                for k, v in obj.items()
            }
        if isinstance(obj, list):
            return [self._mask_nested(item) for item in obj]
        return obj

    @staticmethod
    def _mask_string(value: str) -> str:
        """对字符串进行脱敏，返回统一占位符。"""
        return "***"


# ----------------------------------------------------------------------
# 模块级默认实例与便捷函数（向后兼容现有调用）
# ----------------------------------------------------------------------

_default_helper = CryptoHelper()


def set_crypto_helper(helper: CryptoHelper) -> None:
    """替换全局默认 crypto helper，主要用于测试注入。"""
    global _default_helper
    _default_helper = helper


def reset_crypto_helper() -> None:
    """重置全局默认 crypto helper 为未配置密钥的实例，主要用于测试隔离。"""
    global _default_helper
    _default_helper = CryptoHelper()


def get_crypto_helper() -> CryptoHelper:
    """获取全局默认 crypto helper 实例。"""
    return _default_helper


def validate_encryption_key(environment: str) -> None:
    """校验生产环境是否正确配置加密密钥（委托给默认实例）。"""
    _default_helper.validate_for_environment(environment)


def encrypt_sensitive_fields(config: dict[str, Any] | None) -> dict[str, Any] | None:
    """递归加密配置字典中的敏感字符串字段（委托给默认实例）。"""
    return _default_helper.encrypt_sensitive_fields(config)


def decrypt_sensitive_fields(config: dict[str, Any] | None) -> dict[str, Any] | None:
    """递归解密配置字典中的敏感字符串字段（委托给默认实例）。"""
    return _default_helper.decrypt_sensitive_fields(config)


def encrypt(value: str) -> str:
    """加密单个字符串，返回带 ``enc:`` 前缀的密文（委托给默认实例）。

    密钥未配置时透传原值。注意：加密不幂等，调用方必须保证传入的是明文。
    """
    return _default_helper.encrypt(value)


def decrypt(value: str) -> str:
    """解密单个字符串（委托给默认实例）。

    非 ``enc:`` 前缀的值视为未加密，透传返回；密钥未配置时记录警告并透传；
    解密失败（密钥变更或数据损坏）抛出 ``ValueError``。
    """
    return _default_helper.decrypt(value)


def mask_sensitive_fields(config: dict[str, Any] | None) -> dict[str, Any]:
    """返回配置字典的脱敏副本（委托给默认实例）。"""
    return _default_helper.mask_sensitive_fields(config)


def mask_all_string_values(obj: Any) -> Any:
    """递归将对象中的所有字符串值替换为脱敏占位符（委托给默认实例）。"""
    return _default_helper.mask_all_string_values(obj)