Skip to content

# 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)