1"""Input sanitizer — blocks common prompt-injection patterns before rendering."""
2
3from __future__ import annotations
4
5import re
6
7from lexigram.ai.prompt.exceptions import PromptValidationError
8
9# Patterns that commonly appear in prompt-injection payloads.
10# Each tuple is (name, compiled_pattern).
11_INJECTION_PATTERNS: list[tuple[str, re.Pattern[str]]] = [
12 (
13 "ignore_instructions",
14 re.compile(
15 r"(?i)(ignore|disregard|forget|override)\s+(previous|prior|all|the)\s+"
16 r"(instructions?|prompts?|context|rules?|constraints?)",
17 ),
18 ),
19 (
20 "role_override",
21 re.compile(
22 r"(?i)(you\s+are\s+now|act\s+as|pretend\s+to\s+be|your\s+new\s+role\s+is)",
23 ),
24 ),
25 (
26 "system_prompt_leak",
27 re.compile(
28 r"(?i)(print|show|reveal|output|display|repeat)\s+(your\s+)?"
29 r"(system\s+prompt|instructions?|initial\s+prompt|prompt\s+above)",
30 ),
31 ),
32 (
33 "jinja2_expression",
34 re.compile(r"\{\{.*?\}\}|\{%.*?%\}"),
35 ),
36]
37
38
39class InputSanitizer:
40 """Scans user-supplied variable values for prompt-injection attempts.
41
42 Use :meth:`sanitize` to validate a single string or
43 :meth:`sanitize_all` to check an entire variable mapping.
44
45 Args:
46 strict: When ``True`` (default), detected injections raise
47 :class:`~lexigram.ai.prompt.exceptions.PromptValidationError`.
48 When ``False``, the violating string is returned as-is and a
49 warning is recorded in :attr:`warnings`.
50 extra_patterns: Additional compiled patterns to check.
51 """
52
53 def __init__(
54 self,
55 strict: bool = True,
56 extra_patterns: list[re.Pattern[str]] | None = None,
57 ) -> None:
58 self._strict = strict
59 self._extra: list[re.Pattern[str]] = extra_patterns or []
60 self.warnings: list[str] = []
61
62 def sanitize(self, value: str, variable_name: str = "input") -> str:
63 """Check *value* for injection patterns.
64
65 Args:
66 value: String to inspect.
67 variable_name: Name used in error messages.
68
69 Returns:
70 The (unchanged) *value* when no injection is detected.
71
72 Raises:
73 :class:`~lexigram.ai.prompt.exceptions.PromptValidationError`:
74 Detection in strict mode.
75 """
76 for name, pattern in _INJECTION_PATTERNS:
77 if pattern.search(value):
78 msg = (
79 f"Potential prompt injection detected in '{variable_name}' "
80 f"(pattern: '{name}')."
81 )
82 if self._strict:
83 raise PromptValidationError(msg)
84 self.warnings.append(msg)
85 return value
86
87 for pattern in self._extra:
88 if pattern.search(value):
89 msg = f"Custom injection pattern matched in '{variable_name}'."
90 if self._strict:
91 raise PromptValidationError(msg)
92 self.warnings.append(msg)
93
94 return value
95
96 def sanitize_all(self, variables: dict[str, object]) -> dict[str, object]:
97 """Check all string values in *variables*.
98
99 Non-string values are passed through unchanged.
100
101 Args:
102 variables: Variable name → value mapping.
103
104 Returns:
105 The same mapping (values are not mutated).
106 """
107 for key, value in variables.items():
108 if isinstance(value, str):
109 self.sanitize(value, variable_name=key)
110 return variables
111
112
113__all__ = ["InputSanitizer"]