142 lines
5.5 KiB
Python
142 lines
5.5 KiB
Python
"""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",
|
||
]
|