Secure Your AI Agents on AWS (Part 3): State, Communication, and Detection

# Create an AgentCore memory store with namespace isolation by actor
import boto3

control_client = boto3.client('bedrock-agentcore-control', region_name='us-west-2')
data_client = boto3.client('bedrock-agentcore', region_name='us-west-2')

memory_response = control_client.create_memory(
    name='agent-memory-secure',
    description='Encrypted memory with long-term extraction',
    eventExpiryDuration=90,
    memoryStrategies=[{
        'semanticMemoryStrategy': {
            'name': 'FactExtractor',
            'namespaceTemplates': ['/facts/{actorId}/']
        }
    }]
)

# Same memory store using the AgentCore Python SDK
from bedrock_agentcore.memory import MemoryClient

client = MemoryClient(region_name="us-west-2")

memory = client.create_memory(
    name="agent-memory-secure",
    description="Encrypted memory with long-term extraction",
    strategies=[{
        'semanticMemoryStrategy': {
            'name': 'FactExtractor',
            'namespaceTemplates': ['/facts/{actorId}/']
        }
    }]
)

# Enforce role-based access on RAG retrieval with a metadata filter
response = bedrock_agent_runtime.retrieve(
    knowledgeBaseId='kb-id',
    retrievalQuery={'text': query},
    retrievalConfiguration={
        'vectorSearchConfiguration': {
            # authorized_roles is a list attribute on each document, so use
            # listContains to match when the user's role is one of the allowed
            # roles. Use `equals` only for single-valued scalar attributes.
            'filter': {
                'listContains': {
                    'key': 'authorized_roles',
                    'value': user_role
                }
            }
        }
    }
)

# Create an event bus and allow PutEvents only from tag-authorized agents
import boto3
import json

events = boto3.client('events')
events.create_event_bus(Name='agent-communication-bus')

policy = {
    "Version": "2012-10-17",
    "Statement": [{
        "Sid": "AllowAuthorizedAgents",
        "Effect": "Allow",
        "Principal": {"AWS": "arn:aws:iam::111122223333:root"},
        "Action": "events:PutEvents",
        "Resource": "arn:aws:events:us-west-2:111122223333:event-bus/agent-communication-bus",
        "Condition": {
            "StringEquals": {"aws:PrincipalTag/AgentType": "authorized"}
        }
    }]
}

events.put_permission(
    EventBusName='agent-communication-bus',
    Policy=json.dumps(policy)
)

{
  "Type": "Task",
  "Resource": "arn:aws:states:::aws-sdk:bedrockagentruntime:invokeAgent",
  "Parameters": {
    "AgentId": "<agent-id>",
    "AgentAliasId": "<agent-alias-id>",
    "SessionId.$": "$.sessionId",
    "InputText.$": "$.inputText"
  },
  "Retry": [{
    "ErrorEquals": ["States.TaskFailed"],
    "MaxAttempts": 3,
    "BackoffRate": 2.0
  }],
  "Catch": [{
    "ErrorEquals": ["States.ALL"],
    "ResultPath": "$.error",
    "Next": "IsolateFailure"
  }]
}

# Composite alarm that fires only when error rate AND latency are both high
cloudwatch = boto3.client('cloudwatch')

cloudwatch.put_composite_alarm(
    AlarmName='Agent-CascadingFailure',
    AlarmRule='ALARM(Agent-ErrorRate-High) AND ALARM(Agent-Latency-High)',
    ActionsEnabled=True,
    AlarmActions=['arn:aws:sns:<region>:111122223333:critical-alerts']
)

# Enable GuardDuty with runtime monitoring and malware protection.
# The nested DataSources parameter is deprecated in favor of the flat
# Features list, so configure everything through Features.
import boto3

guardduty = boto3.client('guardduty')

detector = guardduty.create_detector(
    Enable=True,
    FindingPublishingFrequency='FIFTEEN_MINUTES',
    Features=[
        {'Name': 'EKS_AUDIT_LOGS', 'Status': 'ENABLED'},
        {'Name': 'EBS_MALWARE_PROTECTION', 'Status': 'ENABLED'},
        {'Name': 'RUNTIME_MONITORING', 'Status': 'ENABLED'},
        {'Name': 'LAMBDA_NETWORK_LOGS', 'Status': 'ENABLED'}
    ]
)

# Lambda triggered by a Security Hub finding: deny-all the agent role, then notify
import json
import boto3

def contain_rogue_agent(event, context):
    finding = event['detail']['findings'][0]
    # A Security Hub finding resource Id is a full ARN
    # (e.g. arn:aws:iam::111122223333:role/AgentRole-42). IAM APIs take the
    # bare role name, so extract it from the ARN. Splitting on '/' and taking
    # the last segment also handles roles created under a path.
    role_arn = finding['Resources'][0]['Id']
    role_name = role_arn.split('/')[-1]

    # Revoke the agent's credentials immediately with a deny-all policy
    iam = boto3.client('iam')
    iam.put_role_policy(
        RoleName=role_name,
        PolicyName='EmergencyDenyAll',
        PolicyDocument=json.dumps({
            "Version": "2012-10-17",
            "Statement": [{"Effect": "Deny", "Action": "*", "Resource": "*"}]
        })
    )

    # Notify the security team
    sns = boto3.client('sns')
    sns.publish(
        TopicArn='arn:aws:sns:<region>:111122223333:security-critical',
        Subject=f'Rogue Agent Contained: {role_name}',
        Message=f'Agent role {role_name} was automatically contained.'
    )