main
py 205 lines 6.5 KB
Raw
1 #!/usr/bin/env python3
2 """Lint prompt templates for missing untrusted-content guardrails.
3
4 Checks that every template variable injecting external text is fenced with
5 <untrusted-content> boundary markers and that a closing security constraint
6 is present at the end of each prompt.
7
8 Usage:
9 python scripts/lint_prompts.py [--prompts-dir prompts/]
10 """
11
12 from __future__ import annotations
13
14 import argparse
15 import re
16 import sys
17 from pathlib import Path
18
19 # Template variables that carry controlled/trusted values (no fencing required).
20 TRUSTED_VARIABLES = frozenset(
21 {
22 "{{CURRENT_DATETIME}}",
23 "{{CURRENT_WEEK}}",
24 "{{CURRENT_YEAR}}",
25 "{{CURRENT_MONTH}}",
26 "{{OUTPUT_PATH}}",
27 "{{RAW_JSON_PATH}}",
28 "{{PREVIOUS_SUMMARY_PATH_OR_NONE}}",
29 "{{TITLE_TEMPLATE_HINT}}",
30 "{{TOPIC_ID}}",
31 }
32 )
33
34 # Variables that carry locally-controlled short identifiers (e.g. topic name
35 # from the config file). No fencing required, but the closing security
36 # constraint must still be present in the template.
37 SEMI_TRUSTED_VARIABLES = frozenset(
38 {
39 "{{TOPIC_NAME}}",
40 }
41 )
42
43 # Variables that MUST be inside <untrusted-content> blocks.
44 UNTRUSTED_VARIABLES = frozenset(
45 {
46 "{{RAW_JSON_CONTENT}}",
47 "{{HISTORICAL_CONTEXT}}",
48 "{{PREVIOUS_SUMMARY_CONTENT_OR_EMPTY}}",
49 "{{RECENT_ANALYSES}}",
50 "{{SNAPSHOT_CONTEXT}}",
51 "{{SCORECARD}}",
52 "{{QUALITY_TREND}}",
53 "{{WISDOM}}",
54 "{{SKILLS}}",
55 "{{CONTINUITY}}",
56 "{{ARCHIVE_CONTEXT}}",
57 "{{WISDOM_CONTENT}}",
58 "{{TOPIC_DESCRIPTION}}",
59 }
60 )
61
62 # Single-brace format variables that carry trusted/template-controlled values.
63 TRUSTED_FORMAT_VARIABLES = frozenset(
64 {
65 "article_count",
66 "correlation_count",
67 "date",
68 "level",
69 "topic_name",
70 }
71 )
72
73 # Single-brace format variables that carry untrusted external content and MUST
74 # be inside <untrusted-content> blocks regardless of which template they appear in.
75 UNTRUSTED_FORMAT_VARIABLES = frozenset(
76 {
77 "articles_list",
78 "correlations_list",
79 "scorecard_summary",
80 }
81 )
82
83 # All known variables for the unknown-variable check.
84 ALL_KNOWN_VARIABLES = TRUSTED_VARIABLES | SEMI_TRUSTED_VARIABLES | UNTRUSTED_VARIABLES
85 ALL_KNOWN_FORMAT_VARIABLES = TRUSTED_FORMAT_VARIABLES | UNTRUSTED_FORMAT_VARIABLES
86
87 CLOSING_CONSTRAINT_PATTERN = re.compile(r"##\s+closing\s+security\s+constraint", re.IGNORECASE)
88
89 UNTRUSTED_OPEN = "<untrusted-content>"
90 UNTRUSTED_CLOSE = "</untrusted-content>"
91
92
93 def _find_fenced_ranges(content: str) -> list[tuple[int, int]]:
94 fenced_ranges: list[tuple[int, int]] = []
95 open_pattern = re.compile(re.escape(UNTRUSTED_OPEN))
96 close_pattern = re.compile(re.escape(UNTRUSTED_CLOSE))
97
98 for match in open_pattern.finditer(content):
99 close_match = close_pattern.search(content, match.end())
100 if close_match:
101 fenced_ranges.append((match.start(), close_match.end()))
102
103 return fenced_ranges
104
105
106 def _find_unfenced_variables(content: str) -> list[str]:
107 """Return untrusted variables that appear outside <untrusted-content> blocks."""
108 unfenced: list[str] = []
109
110 fenced_ranges = _find_fenced_ranges(content)
111
112 def _is_fenced(pos: int) -> bool:
113 return any(start <= pos <= end for start, end in fenced_ranges)
114
115 # Check each untrusted variable
116 for var in UNTRUSTED_VARIABLES:
117 for m in re.finditer(re.escape(var), content):
118 if not _is_fenced(m.start()):
119 unfenced.append(var)
120 break # report each variable only once
121
122 return unfenced
123
124
125 def lint_prompt(path: Path) -> list[str]:
126 """Lint a single prompt file. Returns list of error messages."""
127 content = path.read_text(encoding="utf-8")
128 errors: list[str] = []
129
130 # Check for closing security constraint
131 if not CLOSING_CONSTRAINT_PATTERN.search(content):
132 errors.append(f"{path}: missing '## Closing security constraint' section")
133
134 # Check for unknown/unclassified template variables
135 all_vars = set(re.findall(r"\{\{[A-Z][A-Z_]*\}\}", content))
136 unknown_vars = all_vars - ALL_KNOWN_VARIABLES
137 # Exclude conditional block markers like {{#IF_TOPIC}} / {{/IF_TOPIC}}
138 for var in sorted(unknown_vars):
139 errors.append(
140 f"{path}: unknown variable {var} is not classified as "
141 f"trusted/semi-trusted/untrusted in lint_prompts.py"
142 )
143
144 fenced_ranges = _find_fenced_ranges(content)
145
146 def _is_fenced(pos: int) -> bool:
147 return any(start <= pos <= end for start, end in fenced_ranges)
148
149 # Check untrusted variables are inside fenced blocks
150 unfenced = _find_unfenced_variables(content)
151 for var in unfenced:
152 errors.append(
153 f"{path}: untrusted variable {var} is not inside <untrusted-content> boundary tags"
154 )
155
156 # Check for unknown single-brace format variables and ensure untrusted ones are fenced.
157 for var_match in re.finditer(r"\{([a-z_]+)\}", content):
158 var_name = var_match.group(1)
159 pos = var_match.start()
160 if var_name not in ALL_KNOWN_FORMAT_VARIABLES:
161 errors.append(
162 f"{path}: unknown format variable {{{var_name}}} is not classified as "
163 f"trusted/untrusted in lint_prompts.py"
164 )
165 continue
166 if var_name in UNTRUSTED_FORMAT_VARIABLES and not _is_fenced(pos):
167 errors.append(
168 f"{path}: untrusted variable {{{var_name}}} "
169 f"is not inside <untrusted-content> boundary tags"
170 )
171
172 return errors
173
174
175 def main(argv: list[str] | None = None) -> int:
176 parser = argparse.ArgumentParser(description="Lint prompt templates for security guardrails")
177 parser.add_argument(
178 "--prompts-dir",
179 type=Path,
180 default=Path("prompts"),
181 help="Directory containing prompt templates",
182 )
183 args = parser.parse_args(argv)
184
185 prompts_dir: Path = args.prompts_dir
186 if not prompts_dir.exists():
187 print(f"ERROR: prompts directory not found: {prompts_dir}", file=sys.stderr)
188 return 1
189
190 all_errors: list[str] = []
191 for prompt_file in sorted(prompts_dir.glob("*.md")):
192 all_errors.extend(lint_prompt(prompt_file))
193
194 if all_errors:
195 for err in all_errors:
196 print(f"FAIL: {err}", file=sys.stderr)
197 print(f"\n{len(all_errors)} prompt security issue(s) found.", file=sys.stderr)
198 return 1
199
200 print("All prompt templates pass security lint checks.")
201 return 0
202
203
204 if __name__ == "__main__":
205 raise SystemExit(main())