utils/crypto.py

"""
对称加密工具（Fernet）：用于把用户自配的 LLM 密钥**加密后**持久化。
=================================================================

用户在「模型设置」面板填写的 API 密钥属于敏感数据。若要跨会话保留（落到本地
SQLite 并经 :mod:`utils.hf_storage` 同步到私有 HF Dataset），**绝不能明文存储**。
本模块用 ``cryptography.fernet.Fernet`` 做对称加密：

- 加密密钥材料来自环境变量 ``PHARMAK_SECRET_KEY``（放 HF Space Secrets）。
  为方便使用，任意长度的字符串都可作为 secret —— 内部用 SHA-256 派生出
  Fernet 所需的 32 字节、URL-safe base64 密钥（故 secret 不必是合法 Fernet key）。
- **优雅降级**：未安装 ``cryptography`` 或未设置 secret 时，``is_available()``
  返回 ``False``，调用方应退回"仅会话存储、不持久化密钥"的行为，避免明文落库。
- ``decrypt`` 在 token 损坏 / secret 轮换导致无法解密时返回 ``None``（不抛出），
  调用方据此当作"无持久配置"处理。

安全说明：密文与 secret 永不进入日志；密文随 DB 同步到**私有** Dataset。
secret 一旦轮换，旧密文将无法解密（视为失效，用户需重新填写密钥）。
"""

from __future__ import annotations

import base64
import hashlib
import logging
import os
from typing import Optional

logger = logging.getLogger(__name__)

#: 加密密钥材料的环境变量名（放 HF Space Secrets）。
ENV_SECRET = "PHARMAK_SECRET_KEY"


def _secret() -> str:
    return (os.environ.get(ENV_SECRET) or "").strip()


def _derive_key(secret: str) -> bytes:
    """从任意 secret 字符串派生 Fernet 所需的 32 字节 URL-safe base64 密钥。"""
    digest = hashlib.sha256(secret.encode("utf-8")).digest()
    return base64.urlsafe_b64encode(digest)


def is_available() -> bool:
    """是否可用加密：``cryptography`` 可导入且已配置 ``PHARMAK_SECRET_KEY``。"""
    if not _secret():
        return False
    try:
        from cryptography.fernet import Fernet  # noqa: F401
    except Exception:
        logger.info("未安装 cryptography，用户密钥将不持久化（退回仅会话存储）。")
        return False
    return True


def _fernet():
    from cryptography.fernet import Fernet

    return Fernet(_derive_key(_secret()))


def encrypt(plaintext: str) -> Optional[str]:
    """加密明文，返回密文 token 字符串。不可用或入参为空时返回 ``None``。"""
    if not plaintext or not is_available():
        return None
    try:
        return _fernet().encrypt(plaintext.encode("utf-8")).decode("ascii")
    except Exception as exc:  # noqa: BLE001 - 加密失败不得影响主流程
        logger.warning("密钥加密失败：%s", exc)
        return None


def decrypt(token: str) -> Optional[str]:
    """解密密文 token，返回明文。无法解密（损坏 / secret 轮换）时返回 ``None``。"""
    if not token or not is_available():
        return None
    try:
        return _fernet().decrypt(token.encode("ascii")).decode("utf-8")
    except Exception:  # noqa: BLE001 - token 失效视为无配置
        logger.info("密钥解密失败（token 失效或 secret 已轮换），按无持久配置处理。")
        return None


__all__ = ["is_available", "encrypt", "decrypt", "ENV_SECRET"]