hive-manage-v3/backend/app/services/key_manager.py

"""Key Manager 服务 — .env 密钥读写、脱敏、批量操作。"""
import logging
import os
import shutil
from pathlib import Path
from typing import Optional

from app.config import settings

logger = logging.getLogger(__name__)

# 可管理的密钥白名单
WHITELIST_FIELDS = [
    "DEEPSEEK_API_KEY",
    "OPENAI_API_KEY",
    "ANTHROPIC_API_KEY",
    "GOOGLE_API_KEY",
    "FEISHU_APP_ID",
    "FEISHU_APP_SECRET",
]

PROFILES_ROOT = Path(settings.hermes_profiles_root)


def _parse_env(content: str) -> dict[str, str]:
    """解析 .env 文件内容为 key-value 字典。"""
    result: dict[str, str] = {}
    for line in content.splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#"):
            continue
        if "=" not in stripped:
            continue
        key, _, val = stripped.partition("=")
        key = key.strip()
        val = val.strip().strip("\"'")
        if key in WHITELIST_FIELDS:
            result[key] = val
    return result


def _mask_value(value: str) -> str:
    """脱敏: 前3字符 + **...** + 后4字符。"""
    if len(value) <= 7:
        return value[:3] + "**...**" + (value[-4:] if len(value) > 3 else "")
    return value[:3] + "**...**" + value[-4:]


def _read_env_file(profile: str) -> dict[str, str]:
    """读取指定 Profile 的 .env 文件中白名单字段。"""
    env_path = PROFILES_ROOT / profile / ".env"
    if not env_path.exists():
        return {}
    try:
        content = env_path.read_text(encoding="utf-8")
        return _parse_env(content)
    except OSError as e:
        logger.warning("读取 %s/.env 失败: %s", profile, e)
        return {}


def _write_env_file(profile: str, updates: dict[str, str]) -> bool:
    """写入指定 Profile 的 .env 文件，操作前自动备份。"""
    env_path = PROFILES_ROOT / profile / ".env"
    if not env_path.exists():
        return False

    # 备份
    try:
        shutil.copy2(env_path, env_path.with_suffix(".env.bak"))
    except OSError as e:
        logger.warning("备份 %s/.env 失败: %s", profile, e)

    try:
        content = env_path.read_text(encoding="utf-8")
        lines = content.splitlines(keepends=True)
        updated_keys = set(updates.keys())
        new_lines: list[str] = []
        found: set[str] = set()
        for line in lines:
            stripped = line.strip()
            if not stripped or stripped.startswith("#"):
                new_lines.append(line)
                continue
            if "=" not in stripped:
                new_lines.append(line)
                continue
            key = stripped.partition("=")[0].strip()
            if key in updated_keys:
                new_lines.append(f"{key}={updates[key]}\n")
                found.add(key)
            else:
                new_lines.append(line)
        # 追加未找到的字段
        for key in updated_keys:
            if key not in found:
                new_lines.append(f"{key}={updates[key]}\n")
        env_path.write_text("".join(new_lines), encoding="utf-8")
        logger.info("已写入 %s/.env (%d 个字段)", profile, len(updates))
        return True
    except OSError as e:
        logger.error("写入 %s/.env 失败: %s", profile, e)
        return False


def list_profiles() -> list[str]:
    """返回所有包含 .env 文件的 Profile 名称列表。"""
    if not PROFILES_ROOT.exists():
        return []
    return sorted(
        d.name
        for d in PROFILES_ROOT.iterdir()
        if d.is_dir() and (d / ".env").exists()
    )


def get_keys(profile_names: Optional[list[str]] = None) -> list[dict]:
    """获取指定 Profile 的密钥列表（脱敏显示）。"""
    profiles = profile_names if profile_names else list_profiles()
    result: list[dict] = []
    for name in profiles:
        fields = _read_env_file(name)
        masked = {k: _mask_value(v) for k, v in fields.items()}
        missing = len(WHITELIST_FIELDS) - len(fields)
        result.append({"profile": name, "fields": masked, "missing": missing})
    return result


def batch_overwrite(
    profile_names: list[str],
    key_name: str,
    new_value: str,
) -> list[dict]:
    """批量覆盖指定 Profile 的某个密钥字段。"""
    results: list[dict] = []
    for name in profile_names:
        success = _write_env_file(name, {key_name: new_value})
        results.append({"profile": name, "success": success})
    return results


def batch_scan(profile_names: Optional[list[str]] = None) -> list[dict]:
    """批量检查缺失的密钥字段。"""
    profiles = profile_names if profile_names else list_profiles()
    results: list[dict] = []
    for name in profiles:
        fields = _read_env_file(name)
        missing = [f for f in WHITELIST_FIELDS if f not in fields]
        results.append({
            "profile": name,
            "missing": missing,
            "missing_count": len(missing),
            "total": len(WHITELIST_FIELDS),
        })
    return results


def get_diff(profile_names: list[str]) -> list[dict]:
    """差异视图: 横向 Profile, 纵向字段, 值状态(有/无/长度)。"""
    profiles = profile_names if profile_names else list_profiles()
    field_map: dict[str, dict[str, str | int]] = {}
    for name in profiles:
        fields = _read_env_file(name)
        for key in WHITELIST_FIELDS:
            if key not in field_map:
                field_map[key] = {}
            if key in fields:
                field_map[key][name] = len(fields[key])
            else:
                field_map[key][name] = 0  # 0 表示缺失
    rows: list[dict] = []
    for key in WHITELIST_FIELDS:
        row: dict[str, str | int] = {"field": key}
        for name in profiles:
            val = field_map.get(key, {}).get(name, 0)
            row[name] = val
        rows.append(row)
    return rows


def export_csv(profile_names: Optional[list[str]] = None) -> str:
    """CSV 导出批量检查结果。"""
    scan = batch_scan(profile_names)
    import csv
    import io
    buf = io.StringIO()
    writer = csv.writer(buf)
    writer.writerow(["Profile", "缺失字段", "缺失数", "总字段"])
    for row in scan:
        writer.writerow([
            row["profile"],
            "; ".join(row["missing"]),
            row["missing_count"],
            row["total"],
        ])
    return buf.getvalue()