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