tools-only AI Agent Security Cheat Sheet
Introduction
AI agents are autonomous systems powered by Large Language Models (LLMs) that can reason, plan, use tools, maintain memory, and take actions to accomplish goals. This expanded capability introduces unique security risks beyond traditional LLM prompt injection. This cheat sheet provides best practices to secure AI agent architectures and minimize attack surfaces.
Key Risks
- Prompt Injection (Direct & Indirect): Malicious instructions injected via user input or external data sources (websites, documents, emails) that hijack agent behavior. (See LLM Prompt Injection Prevention Cheat Sheet)
- Tool Abuse & Privilege Escalation: Agents exploiting overly permissive tools to perform unintended actions or access unauthorized resources.
- Data Exfiltration: Sensitive information leaked through tool calls, API requests, or agent outputs.
- Memory Poisoning: Malicious data persisted in agent memory to influence future sessions or other users.
- Goal Hijacking: Manipulating agent objectives to serve attacker purposes while appearing legitimate.
- Excessive Autonomy: Agents taking high-impact actions without appropriate human oversight.
- Cascading Failures: Compromised agents in multi-agent systems propagating attacks to other agents.
- Denial of Wallet (DoW): Attacks causing excessive API/compute costs through unbounded agent loops.
- Sensitive Data Exposure: PII, credentials, or confidential data inadvertently included in agent context or logs.
Best Practices
1. Tool Security & Least Privilege
- Grant agents the minimum tools required for their specific task.
- Implement per-tool permission scoping (read-only vs. write, specific resources).
- Use separate tool sets for different trust levels (e.g., internal vs. user-facing agents).
- Require explicit tool authorization for sensitive operations.
Bad: Over-permissioned tool configuration
# Dangerous: Agent has unrestricted shell access
tools = [
{
"name": "execute_command",
"description": "Execute any shell command",
"allowed_commands": "*" # No restrictions
}
]
Good: Scoped tool with allowlist
# Safe: Restricted to specific, safe commands
tools = [
{
"name": "file_reader",
"description": "Read files from the reports directory",
"allowed_paths": ["/app/reports/*"],
"allowed_operations": ["read"],
"blocked_patterns": ["*.env", "*.key", "*.pem", "*secret*"]
}
]
Tool Authorization Middleware Example (Python)
from functools import wraps
SENSITIVE_TOOLS = ["send_email", "execute_code", "database_write", "file_delete"]
def require_confirmation(func):
@wraps(func)
async def wrapper(tool_name, params, context):
if tool_name in SENSITIVE_TOOLS:
if not context.get("user_confirmed"):
return {
"status": "pending_confirmation",
"message": f"Action '{tool_name}' requires user approval",
"params": sanitize_for_display(params)
}
return await func(tool_name, params, context)
return wrapper
2. Input Validation & Prompt Injection Defense
- Treat all external data as untrusted (user messages, retrieved documents, API responses, emails).
- Implement input sanitization before including external content in agent context.
- Use delimiters and clear boundaries between instructions and data.
- Apply content filtering for known injection patterns.
- Consider using separate LLM calls to validate/summarize untrusted content.
Please refer to the LLM Prompt Injection Prevention Cheat Sheet for detailed techniques.
3. Memory & Context Security
- Validate and sanitize data before storing in agent memory.
- Implement memory isolation between users/sessions.
- Set memory expiration and size limits.
- Audit memory contents for sensitive data before persistence.
- Use cryptographic integrity checks for long-term memory.
Bad: Unvalidated memory storage
# Dangerous: Storing arbitrary user input in persistent memory
def save_memory(agent, user_message, assistant_response):
agent.memory.add({
"user": user_message, # Could contain injection payload
"assistant": assistant_response,
"timestamp": datetime.now()
})
Good: Validated and isolated memory
import hashlib
from datetime import datetime, timedelta
class SecureAgentMemory:
MAX_MEMORY_ITEMS = 100
MAX_ITEM_LENGTH = 5000
MEMORY_TTL_HOURS = 24
def __init__(self, user_id: str, encryption_key: bytes):
self.user_id = user_id
self.encryption_key = encryption_key
self.memories = []
def add(self, content: str, memory_type: str = "conversation"):
# Validate content
if len(content) > self.MAX_ITEM_LENGTH:
content = content[:self.MAX_ITEM_LENGTH]
# Scan for sensitive data patterns
if self._contains_sensitive_data(content):
content = self._redact_sensitive_data(content)
# Scan for injection patterns
content = self._sanitize_injection_attempts(content)
# Create integrity-checked memory entry
entry = {
"content": content,
"type": memory_type,
"timestamp": datetime.utcnow().isoformat(),
"user_id": self.user_id,
"checksum": self._compute_checksum(content)
}
self.memories.append(entry)
self._enforce_limits()
def get_context(self) -> list:
"""Retrieve valid, non-expired memories."""
valid_memories = []
cutoff = datetime.utcnow() - timedelta(hours=self.MEMORY_TTL_HOURS)
for mem in self.memories:
mem_time = datetime.fromisoformat(mem["timestamp"])
if mem_time > cutoff and self._verify_checksum(mem):
valid_memories.append(mem["content"])
return valid_memories
def _contains_sensitive_data(self, content: str) -> bool:
sensitive_patterns = [
r'\b\d{3}-\d{2}-\d{4}\b', # SSN
r'\b\d{16}\b', # Credit card
r'password\s*[:=]\s*\S+', # Passwords
r'api[_-]?key\s*[:=]\s*\S+' # API keys
]
return any(re.search(p, content, re.I) for p in sensitive_patterns)
def _compute_checksum(self, content: str) -> str:
return hashlib.sha256(
(content + self.user_id).encode() + self.encryption_key
).hexdigest()[:16]
4. Human-in-the-Loop Controls
- Require explicit approval for high-impact or irreversible actions.
- Implement action previews before execution.
- Set autonomy boundaries based on action risk levels.
- Provide clear audit trails of agent decisions and actions.
- Allow users to interrupt and rollback agent operations.
Action Classification and Approval Flow
from enum import Enum
from dataclasses import dataclass
class RiskLevel(Enum):
LOW = "low" # Read operations, safe queries
MEDIUM = "medium" # Write operations, API calls
HIGH = "high" # Financial, deletion, external comms
CRITICAL = "critical" # Irreversible, security-sensitive
ACTION_RISK_MAPPING = {
"search_documents": RiskLevel.LOW,
"read_file": RiskLevel.LOW,
"write_file": RiskLevel.MEDIUM,
"send_email": RiskLevel.HIGH,
"execute_code": RiskLevel.HIGH,
"database_delete": RiskLevel.CRITICAL,
"transfer_funds": RiskLevel.CRITICAL,
}
@dataclass
class PendingAction:
action_id: str
tool_name: str
parameters: dict
risk_level: RiskLevel
explanation: str
class HumanInTheLoopController:
def __init__(self, auto_approve_threshold: RiskLevel = RiskLevel.LOW):
self.auto_approve_threshold = auto_approve_threshold
self.pending_actions = {}
async def request_action(self, tool_name: str, params: dict,
explanation: str) -> dict:
risk_level = ACTION_RISK_MAPPING.get(tool_name, RiskLevel.HIGH)
# Auto-approve low-risk actions
if risk_level.value <= self.auto_approve_threshold.value:
return {"approved": True, "auto": True}
# Queue for human review
action = PendingAction(
action_id=generate_uuid(),
tool_name=tool_name,
parameters=self._sanitize_params_for_display(params),
risk_level=risk_level,
explanation=explanation
)
self.pending_actions[action.action_id] = action
return {
"approved": False,
"pending": True,
"action_id": action.action_id,
"requires": "human_approval",
"risk_level": risk_level.value,
"preview": self._generate_action_preview(action)
}
def _generate_action_preview(self, action: PendingAction) -> str:
return f"""
Action: {action.tool_name}
Risk Level: {action.risk_level.value.upper()}
Explanation: {action.explanation}
Parameters: {json.dumps(action.parameters, indent=2)}
"""
5. Output Validation & Guardrails
- Validate agent outputs before execution or display.
- Implement output filtering for sensitive data leakage.
- Use structured outputs with schema validation where possible.
- Set boundaries on output actions (rate limits, scope limits).
- Apply content safety filters to generated responses.
Output Validation Pipeline
from pydantic import BaseModel, validator
from typing import Optional, List
class AgentToolCall(BaseModel):
tool_name: str
parameters: dict
reasoning: Optional[str]
@validator('tool_name')
def validate_tool_allowed(cls, v):
allowed_tools = ["search", "read_file", "calculator", "get_weather"]
if v not in allowed_tools:
raise ValueError(f"Tool '{v}' is not in allowed list")
return v
@validator('parameters')
def validate_no_sensitive_data(cls, v):
sensitive_patterns = [
r'api[_-]?key', r'password', r'secret', r'token',
r'credential', r'private[_-]?key'
]
params_str = json.dumps(v).lower()
for pattern in sensitive_patterns:
if re.search(pattern, params_str):
raise ValueError("Parameters contain potentially sensitive data")
return v
class OutputGuardrails:
def __init__(self):
self.pii_patterns = self._load_pii_patterns()
self.blocked_actions = set()
self.rate_limiter = RateLimiter(max_calls=100, window_seconds=60)
async def validate_output(self, agent_output: dict) -> dict:
# Check rate limits
if not self.rate_limiter.allow():
raise RateLimitExceeded("Agent action rate limit exceeded")
# Validate structure
if "tool_calls" in agent_output:
for call in agent_output["tool_calls"]:
validated = AgentToolCall(**call)
# Filter PII from responses
if "response" in agent_output:
agent_output["response"] = self._filter_pii(agent_output["response"])
# Check for data exfiltration patterns
if self._detect_exfiltration_attempt(agent_output):
raise SecurityViolation("Potential data exfiltration detected")
return agent_output
def _detect_exfiltration_attempt(self, output: dict) -> bool:
"""Detect attempts to exfiltrate data through tool calls."""
suspicious_patterns = [
# Encoding sensitive data in URLs
lambda o: "http" in str(o) and any(
p in str(o).lower() for p in ["base64", "encode", "password"]
),
# Large data in webhook/API calls
lambda o: o.get("tool_name") in ["http_request", "webhook"] and
len(str(o.get("parameters", ""))) > 10000,
]
return any(pattern(output) for pattern in suspicious_patterns)
6. Monitoring & Observability
- Log all agent decisions, tool calls, and outcomes.
- Implement anomaly detection for unusual agent behavior.
- Track token usage and costs per session/user.
- Set up alerts for security-relevant events.
- Maintain audit trails for compliance and forensics.
Agent Monitoring
import structlog
from dataclasses import dataclass, field
from typing import List, Dict, Any
from datetime import datetime
logger = structlog.get_logger()
@dataclass
class AgentSecurityEvent:
event_type: str
severity: str # INFO, WARNING, CRITICAL
agent_id: str
session_id: str
user_id: str
timestamp: datetime
details: Dict[str, Any]
tool_name: Optional[str] = None
class AgentMonitor:
ANOMALY_THRESHOLDS = {
"tool_calls_per_minute": 30,
"failed_tool_calls": 5,
"injection_attempts": 1,
"sensitive_data_access": 3,
"cost_per_session_usd": 10.0,
}
def __init__(self, agent_id: str):
self.agent_id = agent_id
self.session_metrics = {}
self.alert_handlers = []
async def log_tool_call(self, session_id: str, tool_name: str,
params: dict, result: dict, user_id: str):
# Redact sensitive data before logging
safe_params = self._redact_sensitive(params)
safe_result = self._redact_sensitive(result)
event = AgentSecurityEvent(
event_type="tool_call",
severity="INFO",
agent_id=self.agent_id,
session_id=session_id,
user_id=user_id,
timestamp=datetime.utcnow(),
tool_name=tool_name,
details={
"parameters": safe_params,
"result_status": result.get("status"),
"execution_time_ms": result.get("execution_time_ms"),
}
)
await self._emit_event(event)
await self._check_anomalies(session_id, event)
async def log_security_event(self, session_id: str, event_type: str,
severity: str, details: dict, user_id: str):
event = AgentSecurityEvent(
event_type=event_type,
severity=severity,
agent_id=self.agent_id,
session_id=session_id,
user_id=user_id,
timestamp=datetime.utcnow(),
details=details
)
await self._emit_event(event)
if severity == "CRITICAL":
await self._trigger_alert(event)
async def _check_anomalies(self, session_id: str, event: AgentSecurityEvent):
metrics = self.session_metrics.setdefault(session_id, {
"tool_calls": [],
"failed_calls": 0,
"total_cost": 0.0,
})
metrics["tool_calls"].append(datetime.utcnow())
# Check tool call rate
recent_calls = [t for t in metrics["tool_calls"]
if (datetime.utcnow() - t).seconds < 60]
if len(recent_calls) > self.ANOMALY_THRESHOLDS["tool_calls_per_minute"]:
await self.log_security_event(
session_id, "anomaly_detected", "WARNING",
{"reason": "excessive_tool_calls", "count": len(recent_calls)},
event.user_id
)
def _redact_sensitive(self, data: dict) -> dict:
"""Redact sensitive fields from log data."""
sensitive_keys = {"password", "api_key", "token", "secret", "credential"}
def redact(obj):
if isinstance(obj, dict):
return {
k: "***REDACTED***" if k.lower() in sensitive_keys else redact(v)
for k, v in obj.items()
}
elif isinstance(obj, list):
return [redact(i) for i in obj]
return obj
return redact(data)
7. Multi-Agent Security
- Implement trust boundaries between agents.
- Validate and sanitize inter-agent communications.
- Prevent privilege escalation through agent chains.
- Isolate agent execution environments.
- Apply circuit breakers to prevent cascading failures.
Secure Multi-Agent Communication
from typing import Optional
import jwt
from datetime import datetime, timedelta
class AgentTrustLevel(Enum):
UNTRUSTED = 0
INTERNAL = 1
PRIVILEGED = 2
SYSTEM = 3
class SecureAgentBus:
"""Secure communication layer for multi-agent systems."""
def __init__(self, signing_key: bytes):
self.signing_key = signing_key
self.agent_registry = {}
self.message_validators = []
self.circuit_breakers = {}
def register_agent(self, agent_id: str, trust_level: AgentTrustLevel,
allowed_recipients: List[str]):
self.agent_registry[agent_id] = {
"trust_level": trust_level,
"allowed_recipients": allowed_recipients,
"allowed_message_types": self._get_allowed_types(trust_level)
}
self.circuit_breakers[agent_id] = CircuitBreaker(
failure_threshold=5,
recovery_timeout=60
)
async def send_message(self, sender_id: str, recipient_id: str,
message_type: str, payload: dict) -> dict:
# Validate sender
sender = self.agent_registry.get(sender_id)
if not sender:
raise SecurityViolation(f"Unknown sender agent: {sender_id}")
# Check circuit breaker
if self.circuit_breakers[sender_id].is_open:
raise CircuitBreakerOpen(f"Agent {sender_id} is temporarily blocked")
# Validate recipient authorization
if recipient_id not in sender["allowed_recipients"]:
await self._log_security_event(
"unauthorized_message_attempt",
{"sender": sender_id, "recipient": recipient_id}
)
raise SecurityViolation("Sender not authorized to message recipient")
# Validate message type
if message_type not in sender["allowed_message_types"]:
raise SecurityViolation(f"Message type '{message_type}' not allowed")
# Sanitize payload
sanitized_payload = self._sanitize_payload(payload, sender["trust_level"])
# Create signed message
signed_message = {
"sender": sender_id,
"recipient": recipient_id,
"type": message_type,
"payload": sanitized_payload,
"timestamp": datetime.utcnow().isoformat(),
"signature": self._sign_message(sender_id, recipient_id,
message_type, sanitized_payload)
}
return signed_message
async def receive_message(self, recipient_id: str, message: dict) -> dict:
# Verify signature
if not self._verify_signature(message):
raise SecurityViolation("Invalid message signature")
# Check message freshness (prevent replay attacks)
msg_time = datetime.fromisoformat(message["timestamp"])
if (datetime.utcnow() - msg_time) > timedelta(minutes=5):
raise SecurityViolation("Message expired (possible replay attack)")
# Validate recipient
if message["recipient"] != recipient_id:
raise SecurityViolation("Message recipient mismatch")
return message["payload"]
def _sanitize_payload(self, payload: dict, trust_level: AgentTrustLevel) -> dict:
"""Remove sensitive data based on trust level."""
if trust_level < AgentTrustLevel.PRIVILEGED:
# Remove system-level fields for lower trust agents
payload = {k: v for k, v in payload.items()
if not k.startswith("_system")}
# Always sanitize potential injection content
return sanitize_untrusted_content(payload)
8. Data Protection & Privacy
- Minimize sensitive data in agent context.
- Implement data classification and handling rules.
- Apply encryption for data at rest and in transit.
- Enforce data retention and deletion policies.
- Comply with privacy regulations (GDPR, CCPA).
Data Classification and Handling
from enum import Enum
from typing import Callable
class DataClassification(Enum):
PUBLIC = "public"
INTERNAL = "internal"
CONFIDENTIAL = "confidential"
RESTRICTED = "restricted" # PII, financial, health
class DataProtectionPolicy:
def __init__(self):
self.classification_rules = []
self.handling_rules = {}
def classify_data(self, data: str, context: dict) -> DataClassification:
"""Automatically classify data based on content patterns."""
patterns = {
DataClassification.RESTRICTED: [
r'\b\d{3}-\d{2}-\d{4}\b', # SSN
r'\b\d{16}\b', # Credit card
r'\b[A-Z]{2}\d{6,9}\b', # Passport
r'diagnosis|prescription|patient', # Health
],
DataClassification.CONFIDENTIAL: [
r'salary|compensation|bonus',
r'api[_-]?key|password|secret',
r'confidential|internal only',
],
DataClassification.INTERNAL: [
r'@company\.com',
r'internal|draft|not for distribution',
]
}
for classification, pattern_list in patterns.items():
if any(re.search(p, data, re.I) for p in pattern_list):
return classification
return DataClassification.PUBLIC
def apply_protection(self, data: str, classification: DataClassification,
operation: str) -> str:
"""Apply appropriate protection based on classification."""
handlers = {
DataClassification.RESTRICTED: {
"include_in_context": self._redact_fully,
"log": self._redact_fully,
"output": self._redact_fully,
},
DataClassification.CONFIDENTIAL: {
"include_in_context": self._mask_partially,
"log": self._redact_fully,
"output": self._mask_partially,
},
DataClassification.INTERNAL: {
"include_in_context": lambda x: x,
"log": self._mask_partially,
"output": lambda x: x,
},
}
handler = handlers.get(classification, {}).get(operation, lambda x: x)
return handler(data)
def _redact_fully(self, data: str) -> str:
return "[REDACTED]"
def _mask_partially(self, data: str) -> str:
if len(data) <= 4:
return "****"
return data[:2] + "*" * (len(data) - 4) + data[-2:]
# Usage in agent context building
class SecureContextBuilder:
def __init__(self, policy: DataProtectionPolicy):
self.policy = policy
def build_context(self, documents: List[str], max_tokens: int = 4000) -> str:
protected_docs = []
for doc in documents:
classification = self.policy.classify_data(doc, {})
protected = self.policy.apply_protection(
doc, classification, "include_in_context"
)
protected_docs.append(protected)
# Combine and truncate
context = "\n---\n".join(protected_docs)
return context[:max_tokens * 4] # Rough char estimate
Do's and Don'ts
Do:
- Apply least privilege to all agent tools and permissions.
- Validate and sanitize all external inputs (user messages, documents, API responses).
- Implement human-in-the-loop for high-risk actions.
- Isolate memory and context between users/sessions.
- Monitor agent behavior and set up anomaly detection.
- Use structured outputs with schema validation.
- Sign and verify inter-agent communications.
- Classify data and apply appropriate protections.
Don't:
- Give agents unrestricted tool access or wildcard permissions.
- Trust content from external sources (websites, emails, documents).
- Allow agents to execute arbitrary code without sandboxing.
- Store sensitive data in agent memory without encryption/redaction.
- Let agents make high-impact decisions without human oversight.
- Ignore cost controls (unbounded loops can cause DoW).
- Pass unsanitized data between agents in multi-agent systems.
- Log sensitive data (PII, credentials) in plain text.
References
- OWASP LLM Top 10
- OWASP LLM Prompt Injection Prevention Cheat Sheet
- NIST AI Risk Management Framework
- OpenAI Safety Best Practices
- Google Secure AI Framework (SAIF)