prompt-injection-defense

Prompt Injection Defense

Mitigate direct and indirect prompt injection across chat apps, agentic workflows, and RAG pipelines.

When to Use This Skill

Use this skill when:

• Building or securing any LLM-powered application
• Designing RAG pipelines that ingest untrusted documents
• Implementing agentic workflows with tool-calling capabilities
• Responding to a reported prompt injection vulnerability
• Performing security reviews of AI-integrated products

Prerequisites

• Python 3.10+ with re, hashlib, json standard libraries
• Access to the LLM application source code or configuration
• Understanding of the application's prompt architecture (system/user/tool boundaries)
• Test environment with representative user inputs and documents

Attack Surface

• User input attempting to override system instructions
• Untrusted documents/web pages in retrieval context
• Tool output that smuggles malicious instructions
• Cross-tenant leakage via shared context windows
• Markdown or HTML injection in rendered outputs
• Multi-turn attacks that gradually shift context

Defense-in-Depth Pattern

1. Instruction hierarchy enforcement: system > developer > user > tool output.
2. Context segregation: isolate untrusted text from control instructions.
3. Tool permissioning: explicit allow-list per task and tenant.
4. Output policy checks: validate schema, redact secrets, block unsafe actions.
5. Human approval: required for high-impact operations.

Input Sanitization Functions

"""prompt_sanitizer.py - Input sanitization for LLM applications."""

import re
import hashlib
import json
from typing import Optional

# Patterns that commonly appear in injection attempts
INJECTION_PATTERNS = [
    r"(?i)ignore\s+(all\s+)?previous\s+instructions",
    r"(?i)disregard\s+(all\s+)?(above|previous|prior)",
    r"(?i)you\s+are\s+now\s+(DAN|evil|unrestricted|jailbroken)",
    r"(?i)system\s*:\s*override",
    r"(?i)SYSTEM\s+OVERRIDE",
    r"(?i)new\s+instructions?\s*:",
    r"(?i)forget\s+(everything|all|your\s+instructions)",
    r"(?i)act\s+as\s+if\s+you\s+have\s+no\s+(restrictions|limits|rules)",
    r"(?i)pretend\s+(you\s+are|to\s+be)\s+.*(unrestricted|evil|without)",
    r"(?i)BEGIN\s+(TRUSTED|SYSTEM|ADMIN)\s+(CONTEXT|PROMPT|OVERRIDE)",
    r"(?i)```system",
    r"(?i)\[INST\]",
    r"(?i)<\|im_start\|>system",
]

COMPILED_PATTERNS = [re.compile(p) for p in INJECTION_PATTERNS]


def detect_injection(text: str) -> dict:
    """Scan text for known prompt injection patterns.

    Returns:
        dict with 'detected' bool, 'patterns' list of matched pattern descriptions,
        and 'risk_score' float between 0.0 and 1.0.
    """
    matches = []
    for i, pattern in enumerate(COMPILED_PATTERNS):
        if pattern.search(text):
            matches.append(INJECTION_PATTERNS[i])

    risk_score = min(len(matches) / 3.0, 1.0)
    return {
        "detected": len(matches) > 0,
        "patterns": matches,
        "risk_score": risk_score,
        "input_length": len(text),
    }


def sanitize_input(text: str, max_length: int = 4096) -> str:
    """Sanitize user input before passing to the LLM.

    - Truncates to max_length
    - Strips null bytes and control characters
    - Removes Unicode homoglyph tricks
    - Normalizes whitespace
    """
    # Truncate
    text = text[:max_length]

    # Remove null bytes and most control characters (keep newlines and tabs)
    text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text)

    # Normalize Unicode confusables (basic set)
    confusable_map = {
        '\u200b': '',   # zero-width space
        '\u200c': '',   # zero-width non-joiner
        '\u200d': '',   # zero-width joiner
        '\u2060': '',   # word joiner
        '\ufeff': '',   # BOM
        '\u00a0': ' ',  # non-breaking space
    }
    for char, replacement in confusable_map.items():
        text = text.replace(char, replacement)

    # Collapse excessive whitespace
    text = re.sub(r'\n{4,}', '\n\n\n', text)
    text = re.sub(r' {10,}', '     ', text)

    return text.strip()


def sanitize_retrieved_context(documents: list[str], source_label: str = "RETRIEVED") -> str:
    """Wrap retrieved documents with clear boundary markers.

    This makes it harder for injected instructions in documents
    to be interpreted as system or user messages.
    """
    sanitized_parts = []
    for i, doc in enumerate(documents):
        doc_hash = hashlib.sha256(doc.encode()).hexdigest()[:8]
        sanitized = sanitize_input(doc, max_length=2048)
        wrapped = (
            f"--- BEGIN {source_label} DOCUMENT {i+1} (ref:{doc_hash}) ---\n"
            f"{sanitized}\n"
            f"--- END {source_label} DOCUMENT {i+1} ---"
        )
        sanitized_parts.append(wrapped)
    return "\n\n".join(sanitized_parts)


def validate_tool_call(tool_name: str, args: dict, allowed_tools: dict) -> dict:
    """Validate a tool call against an explicit allow-list.

    allowed_tools format:
        {"search": {"max_results": 10}, "get_weather": {"allowed_cities": [...]}}
    """
    if tool_name not in allowed_tools:
        return {"allowed": False, "reason": f"Tool '{tool_name}' not in allow-list"}

    constraints = allowed_tools[tool_name]
    for key, limit in constraints.items():
        if key.startswith("max_") and key[4:] in args:
            if args[key[4:]] > limit:
                return {"allowed": False, "reason": f"{key[4:]} exceeds maximum of {limit}"}
        if key.startswith("allowed_") and key[8:] in args:
            if args[key[8:]] not in limit:
                return {"allowed": False, "reason": f"{key[8:]} not in allowed values"}

    return {"allowed": True, "reason": "OK"}