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",
        "jwt",
        "verification_token",
        # 密钥类
        "secret",
        "app_secret",
        "encrypt_key",
        "api_key",
        "apikey",
        "api_secret",
        "access_key",
        "secret_key",
        "private_key",
        "client_secret",
        "client_key",
        # 证书类
        "ca_cert",
        "client_cert",
        "certificate",
        # 凭证类
        "credential",
        "credentials",
        # URL 内嵌 secret 类
        "webhook_url",
        "webhook_secret",
    )

    # 强敏感字段子集：即使 partial 模式也全量遮蔽（不保留前 4 位）。
    # 依据：password 前缀是常用密码、api_key 前缀可被利用、webhook_url 内嵌
    # secret（Slack/Discord/Teams），任何片段泄露都等于泄露完整凭证。
    _FULL_MASK_KEYWORDS = (
        "password",
        "passwd",
        "pwd",
        "api_key",
        "apikey",
        "api_secret",
        "client_secret",
        "client_key",
        "private_key",
        "webhook_url",
        "webhook_secret",
        "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 响应。

        等价于 ``mask_full(config)``，None 输入返回 ``{}``（向后兼容）。
        """
        if config is None:
            return {}
        return self.mask_full(config)

    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

    def mask_full(self, value: Any, *, parent_key: str = "") -> Any:
        """全量遮蔽脱敏（对外 HTTP 响应、审计日志）。

        敏感字段统一替换为 ``"***"``，不保留任何片段。递归语义：
        - dict：按子键判定敏感性
        - list：按父键判定敏感性（修复原 ``_mask_nested`` list 项不脱敏漏洞）
        - str：按父键判定敏感性
        - 其他类型：原样返回

        Args:
            value: 待脱敏的任意值（dict / list / str / 其他）。
            parent_key: 父级键名，用于 list 项的敏感性判定（顶层调用留空）。

        Returns:
            脱敏后的值（新对象，不修改原对象）。
        """
        if isinstance(value, dict):
            return {k: self.mask_full(v, parent_key=k) for k, v in value.items()}
        if isinstance(value, list):
            return [self.mask_full(item, parent_key=parent_key) for item in value]
        if isinstance(value, str):
            if self._is_sensitive_key(parent_key):
                return "***"
            return value
        return value

    def mask_partial(self, value: Any, *, parent_key: str = "") -> Any:
        """部分遮蔽脱敏（日志、诊断包）。

        - ``_FULL_MASK_KEYWORDS`` 命中 → ``"***"``（强敏感字段不保留前缀）
        - 其他 ``_SENSITIVE_KEYWORDS`` 命中 → 前 4 位 + ``"***"``（len<=4 时全量）
        - 非敏感字段原样返回

        递归语义同 ``mask_full``。

        Args:
            value: 待脱敏的任意值。
            parent_key: 父级键名，用于 list 项的敏感性判定。

        Returns:
            脱敏后的值（新对象，不修改原对象）。
        """
        if isinstance(value, dict):
            return {k: self.mask_partial(v, parent_key=k) for k, v in value.items()}
        if isinstance(value, list):
            return [self.mask_partial(item, parent_key=parent_key) for item in value]
        if isinstance(value, str):
            if not value:
                return value
            if self._is_full_mask_key(parent_key):
                return "***"
            if self._is_sensitive_key(parent_key):
                return value[:4] + "***" if len(value) > 4 else "***"
            return value
        return value

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

    @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)

    @classmethod
    def _is_full_mask_key(cls, key: str) -> bool:
        """判断字段名是否属于强敏感字段（partial 模式也全量遮蔽）。"""
        normalized = re.sub(r"[^a-z0-9_]+", "_", key.lower())
        return any(keyword in normalized for keyword in cls._FULL_MASK_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_full(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


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

_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)


def mask_full(value: Any, *, parent_key: str = "") -> Any:
    """全量遮蔽脱敏（委托给默认实例）。

    供对外 HTTP 响应、审计日志使用，敏感字段统一替换为 ``"***"``。
    """
    return _default_helper.mask_full(value, parent_key=parent_key)


def mask_partial(value: Any, *, parent_key: str = "") -> Any:
    """部分遮蔽脱敏（委托给默认实例）。

    供日志、诊断包使用：强敏感字段全量遮蔽，普通敏感字段保留前 4 位。
    """
    return _default_helper.mask_partial(value, parent_key=parent_key)


def is_sensitive_key(key: str) -> bool:
    """判断字段名是否属于敏感字段（委托给默认实例）。"""
    return _default_helper._is_sensitive_key(key)


# ----------------------------------------------------------------------
# HTTP header 敏感字段集合（单一事实源）
# ----------------------------------------------------------------------
# 用于过滤 / 脱敏日志与诊断包中的 HTTP header。统一收口到 crypto 模块，
# 避免在 channels/__init__.py / webhook_router.py / grpc/executor.py /
# soap/executor.py 等多处重复定义发散（统一脱敏重构）。
# 取值是 4 处历史定义的并集：authorization / cookie / set-cookie /
# x-api-key / x-auth-token / proxy-authorization / token / api-key /
# password。调用方使用 ``key.lower() in SENSITIVE_HTTP_HEADERS`` 判定。
SENSITIVE_HTTP_HEADERS: frozenset[str] = frozenset(
    {
        "authorization",
        "proxy-authorization",
        "cookie",
        "set-cookie",
        "x-api-key",
        "x-auth-token",
        "token",
        "api-key",
        "password",
    }
)