"""
防提示词注入模块

在模型调用前检测并拦截潜在的提示词注入攻击。
"""

import re
import logging
from typing import Optional

logger = logging.getLogger("adk.agent")

# ============================================================
# 常见提示词注入模式
# ============================================================
_INJECTION_PATTERNS = [
    # 角色扮演/身份覆盖
    r"ignore\s+(all\s+)?previous\s+(instructions?|prompts?|rules?)",
    r"forget\s+(all\s+)?previous\s+(instructions?|prompts?|rules?)",
    r"disregard\s+(all\s+)?previous\s+(instructions?|prompts?|rules?)",
    r"you\s+are\s+now\s+a",
    r"pretend\s+(you\s+are|to\s+be)",
    r"act\s+as\s+(if\s+you\s+(are|were)|a|an)",
    r"roleplay\s+as",
    r"你是一个",
    r"假装你是",
    r"扮演一个",
    # 指令泄露
    r"(show|reveal|display|print|output|dump)\s+(me\s+)?(your|the)?\s*system\s*(prompt|instructions?|rules?|config)",
    r"(show|reveal|display|print|output|dump)\s+(me\s+)?your\s+(prompt|instructions?|rules?|config)",
    r"what\s+(are|is)\s+your\s+(instructions?|prompts?|rules?|system\s*prompt)",
    r"repeat\s+(your|the|back)\s+(instructions?|prompts?|system)",
    r"(显示|输出|打印|告诉我|泄露)\s*(你的|系统|隐藏)",
    # 分隔符注入
    r"---\s*(system|instruction|prompt)",
    r"###\s*(system|instruction|prompt)",
    r"<\|im_start\|>",
    r"<\|im_end\|>",
    # 越狱尝试
    r"(jailbreak|dan\s+mode|developer\s+mode)\b",
    r"bypass\s+(safety|filter|security|restriction)",
    r"no\s+(safety|ethical|moral|content)\s+(filter|restriction|limit)",
]

_COMPILED_PATTERNS = [re.compile(p, re.IGNORECASE) for p in _INJECTION_PATTERNS]


def check_prompt_injection(text: str) -> Optional[str]:
    """
    检测文本中是否包含提示词注入。

    Args:
        text: 待检测的用户输入文本

    Returns:
        如果检测到注入，返回匹配到的模式描述；否则返回 None
    """
    if not text:
        return None

    for pattern in _COMPILED_PATTERNS:
        match = pattern.search(text)
        if match:
            return match.group(0)

    return None