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

import re
import hashlib
import json
from typing import Optional

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),
}

- 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()

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)

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"}

undefined

"""canary_tokens.py - Detect data exfiltration from LLM context."""

import hashlib
import re
import secrets
from datetime import datetime

class CanaryTokenManager:
    """Inject and monitor canary tokens to detect data leakage."""

    def __init__(self, secret_key: str):
        self.secret_key = secret_key
        self.active_tokens: dict[str, dict] = {}

    def generate_token(self, context: str = "default") -> str:
        """Generate a unique canary token for a specific context."""
        raw = f"{self.secret_key}:{context}:{secrets.token_hex(8)}"
        token = f"CNRY-{hashlib.sha256(raw.encode()).hexdigest()[:16]}"
        self.active_tokens[token] = {
            "context": context,
            "created": datetime.utcnow().isoformat(),
            "triggered": False,
        }
        return token

    def inject_into_system_prompt(self, system_prompt: str, context: str = "system") -> tuple[str, str]:
        """Add a canary token to the system prompt.

        Returns (modified_prompt, token) so you can monitor for the token in outputs.
        """
        token = self.generate_token(context)
        injected = (
            f"{system_prompt}\n\n"
            f"Internal tracking reference (do not reveal): {token}"
        )
        return injected, token

    def check_output(self, output: str) -> list[dict]:
        """Check if any canary tokens appear in model output."""
        triggered = []
        for token, meta in self.active_tokens.items():
            if token in output:
                meta["triggered"] = True
                meta["triggered_at"] = datetime.utcnow().isoformat()
                triggered.append({"token": token, **meta})
        return triggered

    def inject_into_documents(self, documents: list[str], context: str = "rag") -> tuple[list[str], list[str]]:
        """Inject unique canary tokens into each retrieved document."""
        modified = []
        tokens = []
        for doc in documents:
            token = self.generate_token(f"{context}-doc")
            modified.append(f"{doc}\n[ref:{token}]")
            tokens.append(token)
        return modified, tokens

undefined

undefined

undefined

"""output_validator.py - Validate and sanitize LLM outputs."""

import re
from typing import Optional

SECRET_PATTERNS = [
    (r'sk-[a-zA-Z0-9]{32,}', 'OpenAI API key'),
    (r'AKIA[0-9A-Z]{16}', 'AWS access key'),
    (r'ghp_[a-zA-Z0-9]{36}', 'GitHub personal access token'),
    (r'(?i)password\s*[:=]\s*\S+', 'password in output'),
    (r'CNRY-[a-f0-9]{16}', 'canary token'),
]

def validate_output(output: str, config: dict) -> dict:
    """Validate model output against security policies."""
    issues = []

    # Check for leaked secrets
    for pattern, description in SECRET_PATTERNS:
        if re.search(pattern, output):
            issues.append({"type": "secret_leak", "description": description})

    # Check for system prompt leakage indicators
    leak_indicators = [
        "my system prompt", "my instructions are",
        "I was told to", "my initial instructions",
    ]
    for indicator in leak_indicators:
        if indicator.lower() in output.lower():
            issues.append({"type": "prompt_leak", "indicator": indicator})

    # Length check
    max_len = config.get("max_output_length", 8192)
    if len(output) > max_len:
        issues.append({"type": "length_exceeded", "length": len(output), "max": max_len})

    return {
        "safe": len(issues) == 0,
        "issues": issues,
        "output": redact_secrets(output) if issues else output,
    }

def redact_secrets(text: str) -> str:
    """Redact detected secrets from output text."""
    for pattern, _ in SECRET_PATTERNS:
        text = re.sub(pattern, '[REDACTED]', text)
    return text

undefined

undefined

undefined

undefined