OWASP Secure Coding Cheat Sheets: AI Agent Security Cheat Sheet

# 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](LLM_Prompt_Injection_Prevention_Cheat_Sheet.md))
- **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.
- **AI Console Malicious Configuration**: AI developer consoles can be compelled to consume data that contains instructions driving malicious changes to the underlying LLM configuration.
- **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.
- **Supply Chain Attacks**: Compromising third-party tools, APIs, or data sources used by agents.

## 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 Model Context Protocol (MCP) tool configuration

```python
# Dangerous: Agent has unrestricted shell access
tools = [
    {
        "name": "execute_command",
        "description": "Execute any shell command",
        "allowed_commands": "*"  # No restrictions
    }
]
```

#### Good: Scoped MCP tool with allowlist

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

```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](LLM_Prompt_Injection_Prevention_Cheat_Sheet.md) 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

```python
# 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

```python
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 has