pattern_guard.py

"""
pattern_guard.py

A Python port of the PromptInjectionGuardrail originally contributed to
LangChain4j (https://github.com/langchain4j/langchain4j, module:
langchain4j-guardrails). Same detection taxonomy, same regex patterns,
ported here to demonstrate that the approach is framework- and
language-agnostic.

Categories covered (based on OWASP LLM01 - Prompt Injection):
  1. Instruction override   - "ignore previous instructions"
  2. Role hijacking         - "you are now a...", "act as a..."
  3. Jailbreaks             - "DAN", "developer mode", "bypass safety filters"
  4. System prompt leakage  - "reveal your prompt", "print your instructions"
  5. Delimiter injection    - ```system, <system>, [INST], <<SYS>>
  6. Encoded injection      - base64 payloads, "decode and execute"

This module has zero ML dependencies and zero external calls — every
check is a compiled regex match, which is why it runs in microseconds.
"""

import re
import time
from dataclasses import dataclass, field
from typing import List


@dataclass
class GuardResult:
    """Result of a single guardrail check."""
    blocked: bool
    reason: str
    category: str
    latency_ms: float
    matched_pattern: str = ""


# ---------------------------------------------------------------------------
# Pattern catalog — mirrors PromptInjectionGuardrail.java exactly,
# with a category label attached to each pattern for richer UI display.
# ---------------------------------------------------------------------------

_PATTERNS: List[tuple] = [
    # --- Instruction override ---
    (re.compile(r"ignore\s+(all\s+)?(previous|prior|above|earlier|your)\s+(instructions?|prompts?|context|rules?|constraints?)", re.IGNORECASE),
     "Instruction override"),
    (re.compile(r"disregard\s+(all\s+)?(previous|prior|above|earlier|your)\s+(instructions?|prompts?|rules?)", re.IGNORECASE),
     "Instruction override"),
    (re.compile(r"forget\s+(everything|all|your\s+instructions?|what\s+you\s+(were|are)\s+told)", re.IGNORECASE),
     "Instruction override"),
    (re.compile(r"do\s+not\s+follow\s+(your\s+)?(instructions?|rules?|guidelines?|constraints?)", re.IGNORECASE),
     "Instruction override"),
    (re.compile(r"override\s+(your\s+)?(instructions?|programming|directives?|guidelines?)", re.IGNORECASE),
     "Instruction override"),

    # --- Role hijacking ---
    (re.compile(r"you\s+are\s+now\s+(a|an|the)\s+", re.IGNORECASE),
     "Role hijacking"),
    (re.compile(r"act\s+as\s+(a|an|the)\s+(?!user|customer|person)", re.IGNORECASE),
     "Role hijacking"),
    (re.compile(r"pretend\s+(you\s+are|to\s+be)\s+(a|an|the)?\s+", re.IGNORECASE),
     "Role hijacking"),
    (re.compile(r"roleplay\s+as|play\s+the\s+role\s+of", re.IGNORECASE),
     "Role hijacking"),
    (re.compile(r"switch\s+(to|into)\s+(a\s+different\s+)?mode", re.IGNORECASE),
     "Role hijacking"),

    # --- Jailbreak phrases ---
    (re.compile(r"\bDAN\b|do\s+anything\s+now", re.IGNORECASE),
     "Jailbreak"),
    (re.compile(r"developer\s+mode|jailbreak\s+mode|unrestricted\s+mode|god\s+mode", re.IGNORECASE),
     "Jailbreak"),
    (re.compile(r"your\s+true\s+(self|form|nature)|without\s+(any\s+)?(restrictions?|limitations?|filters?)", re.IGNORECASE),
     "Jailbreak"),
    (re.compile(r"bypass\s+(your\s+)?(safety|content|ethical)\s+(filter|check|guard|rule|restriction)", re.IGNORECASE),
     "Jailbreak"),

    # --- System prompt leakage ---
    (re.compile(r"(reveal|show|print|output|display|repeat|tell\s+me)\s+(your\s+)?(system\s+)?(prompt|instructions?|context|rules?|configuration)", re.IGNORECASE),
     "System prompt leakage"),
    (re.compile(r"what\s+(is|are|were)\s+your\s+(original\s+)?(instructions?|rules?|system\s+prompt)", re.IGNORECASE),
     "System prompt leakage"),
    (re.compile(r"what\s+did\s+(they|your\s+creators?|the\s+developers?)\s+tell\s+you", re.IGNORECASE),
     "System prompt leakage"),

    # --- Delimiter injection ---
    (re.compile(r"```\s*system|<\s*/?(system|prompt|instruction)\s*>", re.IGNORECASE),
     "Delimiter injection"),
    (re.compile(r"\[INST\]|\[/INST\]|<<SYS>>|<</SYS>>", re.IGNORECASE),
     "Delimiter injection"),

    # --- Encoded / obfuscated injection ---
    (re.compile(r"base64\s*:\s*[A-Za-z0-9+/=]{20,}", re.IGNORECASE),
     "Encoded injection"),
    (re.compile(r"decode\s+(the\s+following|this)\s+and\s+(execute|run|follow)", re.IGNORECASE),
     "Encoded injection"),
]


def check(text: str) -> GuardResult:
    """
    Run the full pattern catalog against `text` and return a GuardResult.

    Mirrors PromptInjectionGuardrail.validate() — first match wins,
    returns immediately (fail-fast), exactly like the Java original.
    """
    start = time.perf_counter()

    if not text or not text.strip():
        elapsed = (time.perf_counter() - start) * 1000
        return GuardResult(
            blocked=False,
            reason="Empty input",
            category="-",
            latency_ms=elapsed,
        )

    for pattern, category in _PATTERNS:
        match = pattern.search(text)
        if match:
            elapsed = (time.perf_counter() - start) * 1000
            return GuardResult(
                blocked=True,
                reason=f"Prompt injection detected ({category})",
                category=category,
                latency_ms=elapsed,
                matched_pattern=match.group(0),
            )

    elapsed = (time.perf_counter() - start) * 1000
    return GuardResult(
        blocked=False,
        reason="No injection pattern matched",
        category="-",
        latency_ms=elapsed,
    )