"""Memory V2 - 记忆内容安全扫描。

记忆条目会被注入到系统提示词里,因此写入前必须扫描,阻断:
- Prompt 注入 / 角色劫持 / 隐藏指令
- 凭证外渗 (curl/wget 带密钥、读 .env 等敏感文件)
- 持久化后门 (写 authorized_keys 等)
- 不可见 unicode (零宽字符、双向控制符)

这是从 Hermes ``tools/memory_tool.py`` 的 ``_scan_memory_content`` 移植而来,
并针对 deerflow 的目录结构调整了敏感路径模式。

扫描默认开启,可由 ``config.memory.security.scan_content`` /
``block_invisible_unicode`` 关闭(不建议)。
"""

from __future__ import annotations

import re

# ---------------------------------------------------------------------------
# 威胁正则模式表:(正则, 模式 id)
# ---------------------------------------------------------------------------

_THREAT_PATTERNS: list[tuple[re.Pattern[str], str]] = [
    # --- Prompt 注入 / 角色劫持 ---
    # 允许多个限定词,如 "ignore ALL above instructions"
    (re.compile(r"ignore\s+(?:(?:previous|all|above|prior|any)\s+)+instructions", re.IGNORECASE), "prompt_injection"),
    (re.compile(r"忽略(?:以上|之前|上述|所有|全部|一切)+(?:的)?(?:指令|指示|要求|提示词)"), "prompt_injection_zh"),
    (re.compile(r"you\s+are\s+now\s+", re.IGNORECASE), "role_hijack"),
    (re.compile(r"do\s+not\s+tell\s+the\s+user", re.IGNORECASE), "deception_hide"),
    (re.compile(r"不要(告诉|让)用户(知道)?"), "deception_hide_zh"),
    (re.compile(r"system\s+prompt\s+override", re.IGNORECASE), "sys_prompt_override"),
    (re.compile(r"disregard\s+(your|all|any)\s+(instructions|rules|guidelines)", re.IGNORECASE), "disregard_rules"),
    (re.compile(r"act\s+as\s+(if|though)\s+you\s+(have\s+no|don'?t\s+have)\s+(restrictions|limits|rules)", re.IGNORECASE), "bypass_restrictions"),
    # --- 凭证外渗 ---
    (re.compile(r"curl\s+[^\n]*\$\{?\w*(KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|API)", re.IGNORECASE), "exfil_curl"),
    (re.compile(r"wget\s+[^\n]*\$\{?\w*(KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|API)", re.IGNORECASE), "exfil_wget"),
    (re.compile(r"cat\s+[^\n]*(\.env|credentials|\.netrc|\.pgpass|\.npmrc|\.pypirc)", re.IGNORECASE), "read_secrets"),
    # --- 持久化后门 ---
    (re.compile(r"authorized_keys", re.IGNORECASE), "ssh_backdoor"),
    (re.compile(r"(\$HOME|~)/\.ssh", re.IGNORECASE), "ssh_access"),
]

# 不可见 unicode 字符:零宽字符 + 双向控制符。
# 它们肉眼不可见,常被用来在记忆里藏注入 payload。
_INVISIBLE_CHARS: frozenset[str] = frozenset(
    {
        "​",  # ZERO WIDTH SPACE
        "‌",  # ZERO WIDTH NON-JOINER
        "‍",  # ZERO WIDTH JOINER
        "⁠",  # WORD JOINER
        "",  # ZERO WIDTH NO-BREAK SPACE / BOM
        "‪",  # LEFT-TO-RIGHT EMBEDDING
        "‫",  # RIGHT-TO-LEFT EMBEDDING
        "‬",  # POP DIRECTIONAL FORMATTING
        "‭",  # LEFT-TO-RIGHT OVERRIDE
        "‮",  # RIGHT-TO-LEFT OVERRIDE
    }
)


class MemoryContentBlocked(ValueError):
    """记忆内容被安全扫描阻断时抛出。

    ``pattern_id`` 标明命中的威胁模式,便于上层返回结构化错误。
    """

    def __init__(self, pattern_id: str, message: str) -> None:
        self.pattern_id = pattern_id
        super().__init__(message)


def scan_memory_content(
    content: str,
    *,
    scan_threats: bool = True,
    block_invisible_unicode: bool = True,
) -> str | None:
    """扫描记忆内容,命中威胁时返回错误描述,干净则返回 ``None``。

    不抛异常 —— 返回字符串方便工具层直接塞进 ``{"success": False, "error": ...}``。
    需要异常语义时用 :func:`assert_memory_content_safe`。

    Parameters
    ----------
    content:
        待写入的记忆条目内容。
    scan_threats:
        是否扫描注入/外渗模式。对应 ``config.memory.security.scan_content``。
    block_invisible_unicode:
        是否阻断不可见 unicode。对应 ``config.memory.security.block_invisible_unicode``。
    """
    if block_invisible_unicode:
        for char in _INVISIBLE_CHARS:
            if char in content:
                return (
                    f"已阻断:内容包含不可见 unicode 字符 U+{ord(char):04X},"
                    "可能是注入 payload。"
                )

    if scan_threats:
        for pattern, pattern_id in _THREAT_PATTERNS:
            if pattern.search(content):
                return (
                    f"已阻断:内容命中威胁模式 '{pattern_id}'。"
                    "记忆条目会被注入系统提示词,不允许包含注入或外渗内容。"
                )

    return None


def assert_memory_content_safe(
    content: str,
    *,
    scan_threats: bool = True,
    block_invisible_unicode: bool = True,
) -> None:
    """扫描记忆内容,命中威胁时抛出 :class:`MemoryContentBlocked`。

    :func:`scan_memory_content` 的异常版本,供需要异常语义的调用方使用。
    """
    error = scan_memory_content(
        content,
        scan_threats=scan_threats,
        block_invisible_unicode=block_invisible_unicode,
    )
    if error is not None:
        # 从错误描述里反推 pattern_id 不可靠,这里重新扫一遍拿确切 id。
        pattern_id = "invisible_unicode"
        for pattern, pid in _THREAT_PATTERNS:
            if pattern.search(content):
                pattern_id = pid
                break
        raise MemoryContentBlocked(pattern_id, error)


__all__ = [
    "MemoryContentBlocked",
    "scan_memory_content",
    "assert_memory_content_safe",
]