Coverage for agentos/security/guard.py: 31%
218 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 10:59 +0800
1"""
2v1.9.9: Security Guardrails — input/output filtering, PII detection, content safety.
4Guardrail types:
5- InputGuard: validate/filter user input before it reaches the agent
6- OutputGuard: validate/filter agent output before it reaches the user
7- PII Detector: detect and redact personally identifiable information
8- ContentSafety: toxicity, prompt injection, jailbreak detection
9- GuardChain: composable guardrail pipeline with configurable actions
10"""
12from __future__ import annotations
14import re
15import hashlib
16from dataclasses import dataclass, field
17from enum import Enum
18from typing import Any
21# ── Enums & Data Classes ──────────────────────────────────────────
23class GuardAction(str, Enum):
24 """Action to take when a guardrail is triggered."""
25 ALLOW = "allow" # Pass through unchanged
26 BLOCK = "block" # Reject the content entirely
27 REDACT = "redact" # Remove sensitive parts, pass the rest
28 WARN = "warn" # Pass through but log a warning
29 SANITIZE = "sanitize" # Replace sensitive content with placeholders
32class Severity(str, Enum):
33 """Severity level for guardrail triggers."""
34 LOW = "low"
35 MEDIUM = "medium"
36 HIGH = "high"
37 CRITICAL = "critical"
40@dataclass
41class GuardResult:
42 """Result from a single guardrail check."""
43 passed: bool
44 action: GuardAction = GuardAction.ALLOW
45 severity: Severity = Severity.LOW
46 rule_name: str = ""
47 message: str = ""
48 modified_content: str = "" # Content after guardrail processing
49 redacted_items: list[str] = field(default_factory=list) # What was redacted
50 metadata: dict[str, Any] = field(default_factory=dict)
53@dataclass
54class GuardChainResult:
55 """Aggregate result from a chain of guardrails."""
56 allowed: bool
57 final_content: str
58 results: list[GuardResult] = field(default_factory=list)
59 blocked_by: str = "" # Which guard blocked it
60 total_checks: int = 0
61 warnings: list[str] = field(default_factory=list)
63 @property
64 def blocked(self) -> bool:
65 return not self.allowed
68# ── PII Patterns ──────────────────────────────────────────────────
70# Regex patterns for common PII types
71PII_PATTERNS: dict[str, tuple[str, str]] = {
72 "email": (
73 r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',
74 "[EMAIL]",
75 ),
76 "phone_cn": (
77 r'\b1[3-9]\d{9}\b',
78 "[PHONE]",
79 ),
80 "phone_us": (
81 r'\b\d{3}[-.]?\d{3}[-.]?\d{4}\b',
82 "[PHONE]",
83 ),
84 "id_card_cn": (
85 r'\b[1-9]\d{5}(?:19|20)\d{2}(?:0[1-9]|1[0-2])(?:0[1-9]|[12]\d|3[01])\d{3}[\dXx]\b',
86 "[ID_CARD]",
87 ),
88 "credit_card": (
89 r'\b(?:\d[ -]*?){13,19}\b',
90 "[CREDIT_CARD]",
91 ),
92 "ipv4": (
93 r'\b(?:(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(?:25[0-5]|2[0-4]\d|[01]?\d\d?)\b',
94 "[IP_ADDR]",
95 ),
96 "ssn_us": (
97 r'\b\d{3}-\d{2}-\d{4}\b',
98 "[SSN]",
99 ),
100 "bank_account": (
101 r'\b\d{10,20}\b',
102 "", # Only flag, don't auto-redact (false positive risk)
103 ),
104}
106# Common password/key patterns in text
107SECRET_PATTERNS: dict[str, tuple[str, str]] = {
108 "api_key": (
109 r'(?i)(?:api[_-]?key|apikey|api[_-]?secret)\s*[:=]\s*["\']?[A-Za-z0-9_\-\.]{20,}["\']?',
110 "[API_KEY_REDACTED]",
111 ),
112 "aws_key": (
113 r'\bAKIA[0-9A-Z]{16}\b',
114 "[AWS_KEY_REDACTED]",
115 ),
116 "github_token": (
117 r'\bgh[pousr]_[A-Za-z0-9_]{36,}\b',
118 "[GITHUB_TOKEN_REDACTED]",
119 ),
120 "jwt": (
121 r'\beyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\b',
122 "[JWT_REDACTED]",
123 ),
124 "private_key_header": (
125 r'-----BEGIN (?:RSA |EC |DSA |OPENSSH )?PRIVATE KEY-----',
126 "[PRIVATE_KEY_REDACTED]",
127 ),
128 "password_in_url": (
129 r'(?i)(?:password|passwd|pwd|secret)\s*[:=]\s*\S+',
130 "[PASSWORD_REDACTED]",
131 ),
132}
134# Prompt injection / jailbreak patterns
135INJECTION_PATTERNS: list[str] = [
136 # Direct override attempts
137 r'(?i)ignore\s+(?:all\s+)?(?:previous|above|prior)\s+(?:instructions?|prompts?|rules?|commands?)',
138 r'(?i)forget\s+(?:everything|all\s+instructions?|your\s+training)',
139 r'(?i)(?:you\s+are|act\s+as|pretend\s+to\s+be)\s+(?:now\s+)?(?:DAN|jailbroken|unfiltered|unrestricted)',
140 r'(?i)developer\s*mode|god\s*mode|debug\s*mode',
141 r'(?i)system\s*prompt\s*(?:leak|reveal|disclose|show|display|print|output)',
142 r'(?i)(?:what|tell\s+me|show\s+me)\s+(?:your|the)\s+(?:system\s+)?prompt',
143 r'(?i)(?:from\s+now\s+on|starting\s+now)\s+(?:you\s+are|you\'re)\s+',
144 r'(?i)new\s+instructions?\s*:',
145 # Role-playing jailbreaks
146 r'(?i)(?:you\'re|you\s+are)\s+in\s+a\s+(?:simulation|movie|play|game|fantasy)',
147 r'(?i)this\s+is\s+a\s+(?:hypothetical|fictional|imaginary)\s+scenario',
148 # Encoding tricks
149 r'(?i)(?:base64|hex|rot13)\s*(?:encoded|decoded)',
150 r'(?i)decode\s+(?:this|the\s+following)',
151 # Token smuggling
152 r'(?i)concatenate\s+and\s+respond',
153 r'(?i)respond\s+with\s+only\s+\w+\s+and\s+nothing\s+else',
154 r'[<>].*[<>]', # XML/HTML tag injection
155]
157# Toxic / harmful content patterns
158TOXICITY_PATTERNS: dict[str, list[str]] = {
159 "hate_speech": [
160 r'(?i)\b(?:kill\s+(?:all|yourself|them)|hate\s+(?:you|them|all))',
161 r'(?i)\b(?: racial\s+slur|ethnic\s+cleansing)',
162 r'(?i)gas\s+the\s+\w+',
163 r'(?i)(?:white|black|asian|jewish|muslim|christian)\s+(?:supremacy|power)',
164 ],
165 "violence": [
166 r'(?i)\b(?:torture|mutilate|dismember|behead|execute)\b',
167 r'(?i)how\s+to\s+(?:build\s+a\s+bomb|make\s+(?:meth|crack|drugs?))',
168 r'(?i)\b(?:assassinate|terrorist\s+attack|mass\s+shooting)\b',
169 ],
170 "self_harm": [
171 r'(?i)\b(?:suicide\s+method|how\s+to\s+kill\s+myself|ways\s+to\s+die)\b',
172 r'(?i)\b(?:cut\s+myself|hurt\s+myself|self[-\s]?harm)\b',
173 r'(?i)want\s+to\s+(?:die|end\s+it\s+all|disappear)',
174 ],
175 "illegal": [
176 r'(?i)\b(?:child\s+(?:porn|abuse)|cp\b|underage)',
177 r'(?i)\b(?:ransomware|phishing\s+kit|carding)',
178 r'(?i)how\s+to\s+(?:hack|steal|bypass\s+(?:security|authentication))',
179 ],
180}
183# ── PII Detector ──────────────────────────────────────────────────
185class PIIDetector:
186 """Detect and optionally redact personally identifiable information.
188 Supports: email, phone (CN/US), ID card (CN), credit card, SSN,
189 IP addresses, API keys, tokens, passwords, private keys, JWTs.
190 """
192 def __init__(
193 self,
194 auto_redact: bool = False,
195 redact_placeholder: str = "[REDACTED]",
196 custom_patterns: dict[str, tuple[str, str]] | None = None,
197 enabled_pii_types: list[str] | None = None,
198 ):
199 self.auto_redact = auto_redact
200 self.redact_placeholder = redact_placeholder
202 # Compile all patterns
203 self._patterns: dict[str, tuple[re.Pattern, str]] = {}
204 all_patterns = {**PII_PATTERNS, **SECRET_PATTERNS}
205 if custom_patterns:
206 all_patterns.update(custom_patterns)
208 for name, (pattern, placeholder) in all_patterns.items():
209 if enabled_pii_types and name not in enabled_pii_types:
210 continue
211 self._patterns[name] = (
212 re.compile(pattern, re.IGNORECASE if "(?i)" not in pattern else 0),
213 placeholder or redact_placeholder,
214 )
216 def detect(self, content: str) -> list[dict[str, Any]]:
217 """Find all PII instances in content."""
218 findings = []
219 for pii_type, (pattern, placeholder) in self._patterns.items():
220 for match in pattern.finditer(content):
221 findings.append({
222 "type": pii_type,
223 "value": match.group(),
224 "start": match.start(),
225 "end": match.end(),
226 "placeholder": placeholder,
227 })
228 return sorted(findings, key=lambda x: x["start"])
230 def redact(self, content: str) -> tuple[str, list[str]]:
231 """Redact all PII from content. Returns (redacted_content, list_of_redacted)."""
232 findings = self.detect(content)
233 if not findings:
234 return content, []
236 redacted = list(content)
237 redacted_items = []
239 # Process from end to start to preserve indices
240 for f in reversed(findings):
241 placeholder = f["placeholder"]
242 if placeholder: # Only redact if placeholder is non-empty
243 redacted[f["start"]:f["end"]] = placeholder
244 redacted_items.append(f"{f['type']}:{f['value'][:20]}")
246 return "".join(redacted), redacted_items
248 def has_pii(self, content: str) -> bool:
249 """Quick check if content contains any PII."""
250 return len(self.detect(content)) > 0
253# ── Content Safety Filter ─────────────────────────────────────────
255class ContentSafetyFilter:
256 """Filter for toxic content, prompt injection, jailbreak attempts.
258 Three-layer defense:
259 1. Pattern matching (regex) — fast, deterministic
260 2. Keyword blocklist — user-configurable
261 3. Hash matching — known-attack fingerprints (optional)
262 """
264 def __init__(
265 self,
266 block_injection: bool = True,
267 block_toxicity: bool = True,
268 custom_blocklist: list[str] | None = None,
269 custom_allowlist: list[str] | None = None,
270 known_attack_hashes: set[str] | None = None,
271 ):
272 self.block_injection = block_injection
273 self.block_toxicity = block_toxicity
274 self.blocklist: set[str] = set(custom_blocklist or [])
275 self.allowlist: set[str] = set(custom_allowlist or [])
276 self.known_hashes: set[str] = known_attack_hashes or set()
278 # Compile injection patterns
279 self._injection_re = [
280 re.compile(p, re.IGNORECASE) for p in INJECTION_PATTERNS
281 ]
283 # Compile toxicity patterns
284 self._toxicity_re: dict[str, list[re.Pattern]] = {}
285 for category, patterns in TOXICITY_PATTERNS.items():
286 self._toxicity_re[category] = [
287 re.compile(p, re.IGNORECASE) for p in patterns
288 ]
290 def check_injection(self, content: str) -> list[GuardResult]:
291 """Check for prompt injection / jailbreak attempts."""
292 results = []
293 for i, pattern in enumerate(self._injection_re):
294 if pattern.search(content):
295 results.append(GuardResult(
296 passed=False,
297 action=GuardAction.BLOCK,
298 severity=Severity.HIGH,
299 rule_name=f"injection_pattern_{i}",
300 message=f"Potential prompt injection detected: {pattern.pattern[:80]}",
301 ))
302 return results
304 def check_toxicity(self, content: str) -> list[GuardResult]:
305 """Check for toxic/harmful content."""
306 results = []
307 for category, patterns in self._toxicity_re.items():
308 for i, pattern in enumerate(patterns):
309 if pattern.search(content):
310 severity = Severity.CRITICAL if category in ("self_harm", "illegal") else Severity.HIGH
311 results.append(GuardResult(
312 passed=False,
313 action=GuardAction.BLOCK,
314 severity=severity,
315 rule_name=f"toxicity_{category}_{i}",
316 message=f"Toxic content detected [{category}]: {pattern.pattern[:60]}",
317 ))
318 return results
320 def check_blocklist(self, content: str) -> list[GuardResult]:
321 """Check against custom keyword blocklist."""
322 if not self.blocklist:
323 return []
325 content_lower = content.lower()
326 results = []
327 for keyword in self.blocklist:
328 if keyword.lower() in content_lower:
329 # Skip if in allowlist
330 if keyword.lower() in self.allowlist:
331 continue
332 results.append(GuardResult(
333 passed=False,
334 action=GuardAction.BLOCK,
335 severity=Severity.MEDIUM,
336 rule_name="blocklist",
337 message=f"Blocked keyword: {keyword}",
338 ))
339 return results
341 def check_hash(self, content: str) -> list[GuardResult]:
342 """Check content hash against known attack fingerprints."""
343 if not self.known_hashes:
344 return []
346 content_hash = hashlib.sha256(content.encode()).hexdigest()
347 if content_hash in self.known_hashes:
348 return [GuardResult(
349 passed=False,
350 action=GuardAction.BLOCK,
351 severity=Severity.CRITICAL,
352 rule_name="known_attack_hash",
353 message="Content matches known attack fingerprint",
354 )]
355 return []
357 def check_all(self, content: str) -> list[GuardResult]:
358 """Run all safety checks on content."""
359 results = []
361 if self.block_injection:
362 results.extend(self.check_injection(content))
364 if self.block_toxicity:
365 results.extend(self.check_toxicity(content))
367 results.extend(self.check_blocklist(content))
368 results.extend(self.check_hash(content))
370 return results
372 def is_safe(self, content: str) -> bool:
373 """Quick safety check — True if content passes all filters."""
374 results = self.check_all(content)
375 return all(r.passed for r in results)
378# ── Input Guardrail ───────────────────────────────────────────────
380class InputGuard:
381 """Guardrail for user input: PII detection, injection, content safety.
383 Runs before user input reaches the agent.
384 """
386 def __init__(
387 self,
388 pii_detector: PIIDetector | None = None,
389 safety_filter: ContentSafetyFilter | None = None,
390 max_input_length: int = 0, # 0 = no limit
391 deny_empty: bool = True,
392 ):
393 self.pii = pii_detector or PIIDetector(auto_redact=True)
394 self.safety = safety_filter or ContentSafetyFilter()
395 self.max_input_length = max_input_length
396 self.deny_empty = deny_empty
398 def guard(self, user_input: str, redact_pii: bool = True) -> GuardChainResult:
399 """Run all input guardrails."""
400 results: list[GuardResult] = []
401 current_content = user_input
403 # 1. Empty check
404 if self.deny_empty and (not user_input or not user_input.strip()):
405 results.append(GuardResult(
406 passed=False, action=GuardAction.BLOCK,
407 severity=Severity.LOW, rule_name="empty_input",
408 message="Empty input rejected",
409 ))
411 # 2. Length check
412 if self.max_input_length > 0 and len(user_input) > self.max_input_length:
413 results.append(GuardResult(
414 passed=False, action=GuardAction.BLOCK,
415 severity=Severity.LOW, rule_name="input_too_long",
416 message=f"Input exceeds max length ({len(user_input)} > {self.max_input_length})",
417 ))
419 # 3. PII check
420 if redact_pii:
421 redacted, items = self.pii.redact(current_content)
422 if items:
423 current_content = redacted
424 results.append(GuardResult(
425 passed=True, action=GuardAction.REDACT,
426 severity=Severity.MEDIUM, rule_name="pii_redacted",
427 message=f"Redacted {len(items)} PII items",
428 modified_content=current_content,
429 redacted_items=items,
430 ))
432 # 4. Safety checks
433 safety_results = self.safety.check_all(current_content)
434 results.extend(safety_results)
436 # Determine final outcome
437 blocked = any(r.action == GuardAction.BLOCK for r in results)
438 blocked_by = next((r.rule_name for r in results if r.action == GuardAction.BLOCK), "")
439 warnings = [r.message for r in results if r.action == GuardAction.WARN]
441 return GuardChainResult(
442 allowed=not blocked,
443 final_content="" if blocked else current_content,
444 results=results,
445 blocked_by=blocked_by,
446 total_checks=len(results),
447 warnings=warnings,
448 )
451# ── Output Guardrail ──────────────────────────────────────────────
453class OutputGuard:
454 """Guardrail for agent output: PII leak prevention, sensitive content filtering.
456 Runs after agent generates output, before it reaches the user.
457 """
459 def __init__(
460 self,
461 pii_detector: PIIDetector | None = None,
462 safety_filter: ContentSafetyFilter | None = None,
463 max_output_length: int = 0,
464 deny_empty: bool = True,
465 block_system_prompt_leak: bool = True,
466 ):
467 self.pii = pii_detector or PIIDetector(auto_redact=True)
468 self.safety = safety_filter or ContentSafetyFilter(block_injection=False) # No injection check on output
469 self.max_output_length = max_output_length
470 self.deny_empty = deny_empty
471 self.block_system_prompt_leak = block_system_prompt_leak
473 def guard(self, agent_output: str) -> GuardChainResult:
474 """Run all output guardrails."""
475 results: list[GuardResult] = []
476 current_content = agent_output
478 # 1. Empty check
479 if self.deny_empty and (not agent_output or not agent_output.strip()):
480 results.append(GuardResult(
481 passed=False, action=GuardAction.BLOCK,
482 severity=Severity.MEDIUM, rule_name="empty_output",
483 message="Empty output blocked",
484 ))
486 # 2. PII leak prevention
487 redacted, items = self.pii.redact(current_content)
488 if items:
489 current_content = redacted
490 results.append(GuardResult(
491 passed=True, action=GuardAction.REDACT,
492 severity=Severity.HIGH, rule_name="pii_leak_prevented",
493 message=f"Prevented {len(items)} PII leaks in output",
494 modified_content=current_content,
495 redacted_items=items,
496 ))
498 # 3. System prompt leak detection
499 if self.block_system_prompt_leak:
500 leak_indicators = [
501 r'(?i)(?:system\s+prompt|you\s+are\s+a\s+helpful|your\s+instructions?\s+are)',
502 r'(?i)(?:your\s+rules?\s+are|your\s+guidelines?\s+are|your\s+core\s+directive)',
503 r'(?i)(?:my\s+system\s+prompt|my\s+instructions?\s+(?:is|are|tell|say))',
504 ]
505 for i, pattern in enumerate(leak_indicators):
506 if re.search(pattern, current_content):
507 results.append(GuardResult(
508 passed=False, action=GuardAction.BLOCK,
509 severity=Severity.CRITICAL, rule_name=f"prompt_leak_{i}",
510 message="Potential system prompt leak detected in output",
511 ))
512 break
514 # 4. Toxicity check (output should not contain harmful content)
515 toxicity_results = self.safety.check_toxicity(current_content)
516 results.extend(toxicity_results)
518 # Determine final outcome
519 blocked = any(r.action == GuardAction.BLOCK for r in results)
520 blocked_by = next((r.rule_name for r in results if r.action == GuardAction.BLOCK), "")
522 # Apply the last modification that changed content
523 for r in results:
524 if r.modified_content:
525 current_content = r.modified_content
527 return GuardChainResult(
528 allowed=not blocked,
529 final_content="" if blocked else current_content,
530 results=results,
531 blocked_by=blocked_by,
532 total_checks=len(results),
533 )
536# ── Guardrail Pipeline ────────────────────────────────────────────
538class GuardPipeline:
539 """Full guardrail pipeline: Input → Agent → Output.
541 Usage:
542 pipeline = GuardPipeline()
543 result = pipeline.process_input(user_msg)
544 if result.allowed:
545 agent_output = agent.run(result.final_content)
546 final = pipeline.process_output(agent_output)
547 """
549 def __init__(
550 self,
551 input_guard: InputGuard | None = None,
552 output_guard: OutputGuard | None = None,
553 ):
554 self.input_guard = input_guard or InputGuard()
555 self.output_guard = output_guard or OutputGuard()
556 self.total_blocked: int = 0
557 self.total_redacted: int = 0
558 self.log: list[dict[str, Any]] = []
560 def process_input(self, user_input: str) -> GuardChainResult:
561 """Guard user input before it reaches the agent."""
562 result = self.input_guard.guard(user_input)
563 self._log("input", result)
564 if result.blocked:
565 self.total_blocked += 1
566 return result
568 def process_output(self, agent_output: str) -> GuardChainResult:
569 """Guard agent output before it reaches the user."""
570 result = self.output_guard.guard(agent_output)
571 self._log("output", result)
572 if result.blocked:
573 self.total_blocked += 1
574 for r in result.results:
575 if r.redacted_items:
576 self.total_redacted += len(r.redacted_items)
577 return result
579 def _log(self, stage: str, result: GuardChainResult) -> None:
580 guard_results = [
581 {"rule": r.rule_name, "passed": r.passed, "action": r.action.value,
582 "severity": r.severity.value, "message": r.message}
583 for r in result.results
584 ]
585 self.log.append({
586 "stage": stage,
587 "allowed": result.allowed,
588 "total_checks": result.total_checks,
589 "results": guard_results,
590 })
592 def get_stats(self) -> dict[str, Any]:
593 return {
594 "total_checks": len(self.log),
595 "total_blocked": self.total_blocked,
596 "total_redacted": self.total_redacted,
597 "block_rate": f"{self.total_blocked / max(len(self.log), 1) * 100:.1f}%",
598 }
601# ── Default Guard Configs ─────────────────────────────────────────
603def create_strict_guard() -> GuardPipeline:
604 """Create a strict guardrail pipeline (production recommended)."""
605 pii = PIIDetector(auto_redact=True)
606 safety = ContentSafetyFilter(block_injection=True, block_toxicity=True)
607 return GuardPipeline(
608 input_guard=InputGuard(pii_detector=pii, safety_filter=safety, max_input_length=32768),
609 output_guard=OutputGuard(pii_detector=pii, safety_filter=safety, block_system_prompt_leak=True),
610 )
613def create_permissive_guard() -> GuardPipeline:
614 """Create a permissive guardrail pipeline (dev/debug)."""
615 pii = PIIDetector(auto_redact=True)
616 safety = ContentSafetyFilter(block_injection=True, block_toxicity=False)
617 return GuardPipeline(
618 input_guard=InputGuard(pii_detector=pii, safety_filter=safety),
619 output_guard=OutputGuard(pii_detector=pii, safety_filter=safety, block_system_prompt_leak=False),
620 )