# -*- coding: utf-8 -*-
"""
敏感信息脱敏过滤器

借鉴 zhikuncode-main `security/SensitiveDataFilter.java`，对工具输出、记忆库写入前的
文本进行 16 模式正则匹配，把泄露的密钥替换为 `前缀-***REDACTED-长度***` 格式，
保留前缀 + 长度信息便于调试识别，但不泄露密钥本身。

集成点：
1. `tools/registry.py:format_tool_result` 出口（保护 LLM 上下文）
2. `memory/store.py:add()` 持久化前（保护跨会话记忆库）
3. `tools/shell_exec.py` / `web_fetch.py` 结果出口（纵深防御）
"""
import re
import threading
from typing import Any, Dict, List, Optional, Tuple


# 16 模式：覆盖主流密钥类型
# 每项：(compiled_pattern, label, min_length)
# min_length 用于过滤过短的误匹配
_PATTERNS: List[Tuple[str, str, int]] = [
    # 1. OpenAI 标准 key（sk- + 40 个 base62 字符）
    (r"sk-([a-zA-Z0-9]{40,})", "OPENAI_KEY", 48),
    # 2. OpenAI project key（sk-proj- + 40 个 base62/-/_ 字符）
    (r"sk-proj-([a-zA-Z0-9\-_]{40,})", "OPENAI_PROJ_KEY", 49),
    # 3. Anthropic key（sk-ant- + 40 个 base62/-/_ 字符）
    (r"sk-ant-([a-zA-Z0-9\-_]{40,})", "ANTHROPIC_KEY", 47),
    # 4. AWS access key id（AKIA + 16 个 base32 字符）
    (r"(AKIA[0-9A-Z]{16})", "AWS_AKIA", 20),
    # 5. GitHub personal access token（ghp_ + 36 个 base62 字符）
    (r"(ghp_[a-zA-Z0-9]{36})", "GITHUB_TOKEN", 40),
    # 6. GitLab personal access token（glpat- + 20 个 base62/-/_ 字符）
    (r"(glpat-[a-zA-Z0-9\-_]{20})", "GITLAB_TOKEN", 26),
    # 7. Slack token（xox[baprs]- + 10+ 字符 + 可选 - + 10+ 字符）
    (r"(xox[baprs]-[a-zA-Z0-9\-]{10,}(?:-[a-zA-Z0-9\-]{10,})?)", "SLACK_TOKEN", 15),
    # 8. HuggingFace token（hf_ + 30 个 base62 字符）
    (r"(hf_[a-zA-Z0-9]{30,})", "HUGGINGFACE_TOKEN", 33),
    # 9. GCP API key（AIza + 35 个 base62/-/_ 字符）
    (r"(AIza[0-9A-Za-z\-_]{35})", "GCP_API_KEY", 39),
    # 10. Stripe live secret key（sk_live_ + 24+ 字符）
    (r"(sk_live_[a-zA-Z0-9]{24,})", "STRIPE_LIVE_KEY", 32),
    # 11. Vercel token（vercel_ + 24+ 字符）— 前缀 7 字符 + 24 字符 = 31
    (r"(vercel_[a-zA-Z0-9]{24,})", "VERCEL_TOKEN", 31),
    # 12. Supabase token（sbp_ + 24+ 字符）
    (r"(sbp_[a-zA-Z0-9]{24,})", "SUPABASE_TOKEN", 28),
    # 13. JWT（三段式，以 eyJ 开头）
    (r"(eyJ[a-zA-Z0-9_-]+\.eyJ[a-zA-Z0-9_-]+\.[a-zA-Z0-9_-]+)", "JWT", 16),
    # 14. PEM 私钥头（含完整 BEGIN/END 块）
    (r"-----BEGIN (?:RSA |EC |OPENSSH |PGP |DSA |)PRIVATE KEY-----.*?-----END (?:RSA |EC |OPENSSH |PGP |DSA |)PRIVATE KEY-----",
     "PEM_PRIVATE_KEY", 50),
    # 15. 通用 key=value 凭证（密码、token、secret 等）
    # 注意：只匹配引号包裹的值，避免误伤
    (r'((?:api[_-]?key|access[_-]?token|secret[_-]?key|secret|password|passwd|pwd|client[_-]?secret)["\']?\s*[:=]\s*["\'])([^\s"\']{8,})(["\'])',
     "GENERIC_CREDENTIAL", 8),
    # 16. 数据库连接串密码（postgres/mysql/mongodb://user:pass@host）
    (r"((?:postgres|postgresql|mysql|mongodb|redis)://[^:\s]+:)([^\s@]+)(@)", "DB_PASSWORD", 1),
]


class SensitiveDataFilter:
    """敏感信息脱敏过滤器。

    单例使用推荐通过 `get_default_filter()`，避免重复编译正则。
    线程安全：内部正则在编译后只读，多线程并发 redact 无需加锁。
    """

    def __init__(self):
        # 编译所有正则，提高重复调用性能
        self._compiled: List[Tuple[re.Pattern, str, int]] = [
            (re.compile(p, re.DOTALL if "PRIVATE KEY" in p else 0), label, min_len)
            for p, label, min_len in _PATTERNS
        ]

    def redact(self, text: str) -> str:
        """对文本中的敏感信息脱敏。

        脱敏格式：`<前缀>***REDACTED-<长度>***`
        - 保留前缀（如 sk-/ghp_/AKIA）便于调试识别密钥类型
        - 标注原密钥长度便于审计
        - 不泄露密钥本身

        Args:
            text: 待处理的文本

        Returns:
            脱敏后的文本（无匹配则原样返回）
        """
        if not text or not isinstance(text, str):
            return text

        result = text
        for pattern, label, min_len in self._compiled:
            result = self._apply_pattern(result, pattern, label, min_len)
        return result

    def _apply_pattern(self, text: str, pattern: re.Pattern, label: str, min_len: int) -> str:
        """对单个模式应用脱敏"""
        def replace_func(match):
            full_match = match.group(0)
            # 长度过短的匹配跳过（避免误判）
            if len(full_match) < min_len:
                return full_match

            # 提取前缀（用于识别密钥类型）
            prefix = self._extract_prefix(full_match, label)

            # 通用凭证模式：group(2) 是密码值
            if label == "GENERIC_CREDENTIAL":
                # group(1) 是 key= 部分，group(2) 是值，group(3) 是结束引号
                key_part = match.group(1)
                value = match.group(2)
                end_quote = match.group(3)
                return f"{key_part}***REDACTED-{len(value)}***{end_quote}"

            # DB 连接串：group(1) 是协议+用户:，group(2) 是密码，group(3) 是@
            if label == "DB_PASSWORD":
                proto_user = match.group(1)
                password = match.group(2)
                at_sign = match.group(3)
                return f"{proto_user}***REDACTED-{len(password)}***{at_sign}"

            # PEM 私钥：整块替换
            if label == "PEM_PRIVATE_KEY":
                return f"***REDACTED-PEM-PRIVATE-KEY***"

            # 其他模式：保留前缀 + 标注长度
            return f"{prefix}***REDACTED-{label}-{len(full_match)}***"

        return pattern.sub(replace_func, text)

    @staticmethod
    def _extract_prefix(match: str, label: str) -> str:
        """从匹配的密钥中提取前缀（用于脱敏后识别密钥类型）"""
        # 找到第一个非字母数字/破折号/下划线的位置
        for i, c in enumerate(match):
            if i > 0 and not (c.isalnum() or c in "-_"):
                return match[:i]
        # 全部是字母数字/破折号/下划线，返回完整（如 JWT）
        if len(match) > 10:
            return match[:6]  # 截取前 6 字符作前缀
        return match

    def redact_dict(self, data: Any) -> Any:
        """递归对 dict / list 中的字符串值脱敏。

        非字符串值（int/float/bool/None）原样保留。
        用于处理工具结果 dict（如 shell_exec 返回的 {stdout, stderr, ...}）。

        Args:
            data: 任意 JSON 兼容数据

        Returns:
            脱敏后的数据（结构相同）
        """
        if isinstance(data, str):
            return self.redact(data)
        if isinstance(data, dict):
            return {k: self.redact_dict(v) for k, v in data.items()}
        if isinstance(data, list):
            return [self.redact_dict(item) for item in data]
        return data


# 模块级单例：避免每次调用都重新编译正则
_default_filter: Optional[SensitiveDataFilter] = None
_default_lock = threading.Lock()


def get_default_filter() -> SensitiveDataFilter:
    """获取默认单例 SensitiveDataFilter。

    首次调用时创建实例并编译所有正则，后续调用直接返回缓存实例。
    线程安全。
    """
    global _default_filter
    if _default_filter is None:
        with _default_lock:
            if _default_filter is None:
                _default_filter = SensitiveDataFilter()
    return _default_filter
