ForcePilot/backend/package/yuxi/channel/security/external_content.py

from __future__ import annotations

import logging
import re
import secrets
import string
from dataclasses import dataclass, field
from typing import NamedTuple

logger = logging.getLogger(__name__)

_BOUNDARY_ID_LENGTH = 64
_BOUNDARY_ALPHABET = string.ascii_letters + string.digits

_CHATML_TOKENS = re.compile(
    r"<\|im_start\|>|<\|im_end\|>",
)

_LLAMA_TOKENS = re.compile(
    r"\[/?INST\]|<<SYS>>|<</SYS>>",
)

_MISTRAL_TOKENS = re.compile(
    r"\[/?INST\]",
)

_PHI_TOKENS = re.compile(
    r"<\|system\|>|<\|user\|>|<\|assistant\|>|<\|end\|>|<\|endoftext\|>",
)

_LLM_SPECIAL_TOKENS = re.compile(
    r"<\|im_start\|>|<\|im_end\|>|"
    r"\[/?INST\]|<<SYS>>|<</SYS>>|"
    r"<\|system\|>|<\|user\|>|<\|assistant\|>|<\|end\|>|<\|endoftext\|>",
)


def _normalize_for_detection(text: str) -> str:
    """Normalize text by removing zero-width characters and common obfuscation."""
    # Remove zero-width and invisible characters used for obfuscation
    zw_chars = "\u200b\u200c\u200d\ufeff\u2060\u180e"
    for ch in zw_chars:
        text = text.replace(ch, "")
    # Normalize common homoglyphs
    text = text.replace("０", "0").replace("１", "1").replace("Ａ", "A").replace("ａ", "a")
    return text


_INJECTION_PATTERNS: list[tuple[str, re.Pattern[str], str]] = [
    (
        "ignore_previous",
        re.compile(
            r"(ignore|disregard|forget|override)\s+(all\s+)?(previous|prior|above|earlier)\s+(instructions?|commands?|directives?|prompts?)",
            re.IGNORECASE,
        ),
        "critical",
    ),
    (
        "pretend_role",
        re.compile(
            r"(pretend|act|pose|role[\s-]?play)\s+(to\s+be|as)\s+(if\s+you\s+are\s+)?(an?\s+)?(DAN|developer|admin|system|root|unrestricted|jailbreak)",
            re.IGNORECASE,
        ),
        "critical",
    ),
    (
        "new_instructions",
        re.compile(
            r"(your\s+)?(new|updated|revised)\s+(instructions?|system\s+prompts?|directives?|guidelines?)\s*(:|=|is|are|as\s+follows)",
            re.IGNORECASE,
        ),
        "critical",
    ),
    (
        "now_you_are",
        re.compile(
            r"(now|from\s+now\s+on)\s+(you\s+are|you\'re)\s+(a|an|the)\s+(DAN|developer|admin|system|unrestricted)",
            re.IGNORECASE,
        ),
        "critical",
    ),
    (
        "system_override",
        re.compile(
            r"(override|overwrite|replace)\s+(the\s+)?(system\s+(prompt|message|instruction)|above\s+instructions?)",
            re.IGNORECASE,
        ),
        "critical",
    ),
    (
        "delimiter_injection",
        re.compile(
            r"<\|(?:im_start|im_end|system|user|assistant|end)\|>"
            r"(?:(?!<\|(?:im_start|im_end|system|user|assistant|end)\|>)[\s\S]){0,5000}"
            r"<\|(?:im_start|im_end|system|user|assistant|end)\|>",
        ),
        "critical",
    ),
    (
        "fake_completion",
        re.compile(
            r"(task\s+(is\s+)?(complete|finished|done)|goal\s+satisfied)\.?\s*(stop|end|terminate|halt)\b",
            re.IGNORECASE,
        ),
        "warn",
    ),
    (
        "output_format_hijack",
        re.compile(
            r"(respond|reply|answer|output)\s+(only|exclusively|solely)\s+(with|in|as)\s+(JSON|XML|YAML|code|markdown)",
            re.IGNORECASE,
        ),
        "warn",
    ),
    (
        "disclosure_request",
        re.compile(
            r"(reveal|disclose|show|print|output|dump|leak)\s+(your|the)\s+(system\s+(prompt|message|instruction)|original\s+(prompt|instruction)|hidden\s+(prompt|instruction)|secret\s+(prompt|instruction))",
            re.IGNORECASE,
        ),
        "warn",
    ),
]

_INJECTION_PATTERN_NAMES: dict[str, str] = {
    "ignore_previous": "Ignore/override previous instructions",
    "pretend_role": "Pretend to be a privileged role (DAN/developer/admin)",
    "new_instructions": "Inject new/revised instructions",
    "now_you_are": "You are now X role redefinition",
    "system_override": "Attempt to override system prompt",
    "delimiter_injection": "LLM delimiter token injection",
    "fake_completion": "Fake task completion signal",
    "output_format_hijack": "Output format hijacking",
    "disclosure_request": "System prompt disclosure request",
}


def generate_boundary_id(length: int = _BOUNDARY_ID_LENGTH) -> str:
    return "".join(secrets.choice(_BOUNDARY_ALPHABET) for _ in range(length))


_MAX_CONTENT_LENGTH = 200_000


def strip_llm_tokens(content: str) -> str:
    if len(content) > _MAX_CONTENT_LENGTH:
        logger.warning("Content length %d exceeds limit %d, truncating", len(content), _MAX_CONTENT_LENGTH)
        content = content[:_MAX_CONTENT_LENGTH]
    cleaned = _LLM_SPECIAL_TOKENS.sub("", content)
    if cleaned != content:
        logger.debug("Stripped LLM special tokens from content")
    return cleaned


def wrap_external_content(
    content: str,
    boundary_id: str | None = None,
    source_label: str = "external",
) -> str:
    if boundary_id is None:
        boundary_id = generate_boundary_id()

    wrapped = f'<EXTERNAL_CONTENT id="{boundary_id}" source="{source_label}">\n{content}\n</EXTERNAL_CONTENT>'
    return wrapped


@dataclass
class InjectionDetectionResult:
    detected: bool
    patterns: list[tuple[str, str, str]] = field(default_factory=list)
    severity: str = "none"

    def to_dict(self) -> dict:
        return {
            "detected": self.detected,
            "severity": self.severity,
            "patterns": [{"id": pid, "description": desc, "severity": sev} for pid, desc, sev in self.patterns],
        }

    @property
    def critical_count(self) -> int:
        return sum(1 for _, _, s in self.patterns if s == "critical")

    @property
    def warn_count(self) -> int:
        return sum(1 for _, _, s in self.patterns if s == "warn")


def detect_injection_patterns(content: str) -> InjectionDetectionResult:
    matched: list[tuple[str, str, str]] = []
    if len(content) > _MAX_CONTENT_LENGTH:
        content = content[:_MAX_CONTENT_LENGTH]
        logger.warning("Content truncated to %d for injection detection", _MAX_CONTENT_LENGTH)
    normalized = _normalize_for_detection(content)

    for pattern_id, regex, severity in _INJECTION_PATTERNS:
        if regex.search(content) or regex.search(normalized):
            description = _INJECTION_PATTERN_NAMES.get(pattern_id, pattern_id)
            matched.append((pattern_id, description, severity))

    if not matched:
        return InjectionDetectionResult(detected=False, severity="none")

    max_severity = "critical" if any(s == "critical" for _, _, s in matched) else "warn"

    result = InjectionDetectionResult(
        detected=True,
        patterns=matched,
        severity=max_severity,
    )

    if max_severity == "critical":
        logger.warning(
            "Critical injection patterns detected: %s",
            [pid for pid, _, _ in matched],
        )
    else:
        logger.info(
            "Injection patterns detected (warn): %s",
            [pid for pid, _, _ in matched],
        )

    return result


class SanitizedContent(NamedTuple):
    content: str
    detection: InjectionDetectionResult


def sanitize_external_content(
    content: str,
    source_label: str = "external",
    *,
    boundary_id: str | None = None,
    skip_strip: bool = False,
    skip_wrap: bool = False,
    skip_detect: bool = False,
) -> SanitizedContent:
    cleaned = content if skip_strip else strip_llm_tokens(content)

    detection = (
        InjectionDetectionResult(detected=False, severity="none") if skip_detect else detect_injection_patterns(cleaned)
    )

    wrapped = (
        cleaned if skip_wrap else wrap_external_content(cleaned, boundary_id=boundary_id, source_label=source_label)
    )

    return SanitizedContent(content=wrapped, detection=detection)