Insecure Inter-Agent Communication: When Agents Talk, Attackers Listen (ASI07)

import boto3
import json
from datetime import datetime, timezone
from jsonschema import validate, FormatChecker, ValidationError

eventbridge = boto3.client("events")
schemas = boto3.client("schemas")

# Define the contract for inter-agent messages
delegation_schema = {
    "openapi": "3.0.0",
    "info": {"title": "AgentDelegation", "version": "1.0.0"},
    "paths": {},
    "components": {
        "schemas": {
            "DelegationMessage": {
                "type": "object",
                "required": ["source_agent", "target_agent", "task_type", "parameters", "nonce", "timestamp"],
                "properties": {
                    "source_agent": {"type": "string", "pattern": "^agent-[a-z]+-[0-9]+$"},
                    "target_agent": {"type": "string", "pattern": "^agent-[a-z]+-[0-9]+$"},
                    "task_type": {"type": "string", "enum": ["query", "process", "notify"]},
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "query": {"type": "string", "maxLength": 500},
                            "resource_id": {"type": "string", "pattern": "^[a-zA-Z0-9-]+$"}
                        },
                        "additionalProperties": False  # No free-form injection fields
                    },
                    "nonce": {"type": "string", "minLength": 32},
                    "timestamp": {"type": "string", "format": "date-time"}
                },
                "additionalProperties": False
            }
        }
    }
}

# Register the schema (documentation + code generation).
# Note: this is a one-time bootstrap. create_registry and create_schema both
# throw if the thing already exists, so do not run this on every invocation.
schemas.create_registry(RegistryName="agent-communications")
schemas.create_schema(
    RegistryName="agent-communications",
    SchemaName="AgentDelegation",
    Type="OpenApi3",
    Content=json.dumps(delegation_schema)
)

# Extract the JSON Schema portion for validation
message_schema = delegation_schema["components"]["schemas"]["DelegationMessage"]

def publish_agent_message(message: dict):
    """Validate against the schema before putting anything on the bus."""
    try:
        validate(instance=message, schema=message_schema, format_checker=FormatChecker())
    except ValidationError as e:
        raise ValueError(f"Message rejected at producer: {e.message}")

    eventbridge.put_events(Entries=[{
        "Source": f"agent.{message['source_agent']}",
        "DetailType": "AgentDelegation",
        "Detail": json.dumps(message),
        "EventBusName": "agent-bus"
    }])

def handle_agent_message(event: dict):
    """Consumer validates independently. Never trust that the producer did its job."""
    # EventBridge delivers `detail` to a Lambda target already parsed into a dict.
    # Do NOT json.loads() it. (If you consume via SQS or Kinesis instead, parse the
    # record body, and pull `detail` out of that.)
    message = event["detail"]

    # Schema validation. Pass a FormatChecker so `date-time` is actually asserted.
    # That check needs the extras: pip install "jsonschema[format]"
    try:
        validate(instance=message, schema=message_schema, format_checker=FormatChecker())
    except ValidationError as e:
        print(f"[REJECTED] Invalid message from {message.get('source_agent', 'unknown')}: {e.message}")
        return  # Drop it. Don't act on it.

    # Anti-replay: check the nonce has not been used before
    if nonce_store.exists(message["nonce"]):
        print(f"[REJECTED] Replay detected: {message['nonce']}")
        return
    nonce_store.add(message["nonce"], ttl=300)

    # Freshness: parse defensively, and compare timezone-aware to timezone-aware
    try:
        msg_time = datetime.fromisoformat(message["timestamp"].replace("Z", "+00:00"))
    except ValueError:
        print(f"[REJECTED] Unparseable timestamp: {message['timestamp']}")
        return
    if msg_time.tzinfo is None:
        msg_time = msg_time.replace(tzinfo=timezone.utc)
    if (datetime.now(timezone.utc) - msg_time).total_seconds() > 60:
        print(f"[REJECTED] Stale message: {message['timestamp']}")
        return

    # Only now do we trust it enough to act
    process_delegation(message)

{
  "Comment": "Secure multi-agent orchestration",
  "StartAt": "ValidateRequest",
  "States": {
    "ValidateRequest": {
      "Type": "Task",
      "Resource": "arn:aws:lambda:us-east-1:123456789012:function:validate-agent-message",
      "Next": "RouteToAgent",
      "Catch": [
        { "ErrorEquals": ["ValidationError"], "ResultPath": "$.error", "Next": "RejectAndLog" },
        { "ErrorEquals": ["States.ALL"],      "ResultPath": "$.error", "Next": "RejectAndLog" }
      ]
    },
    "RouteToAgent": {
      "Type": "Choice",
      "Choices": [
        { "Variable": "$.task_type", "StringEquals": "process_payment", "Next": "PaymentAgent" },
        { "Variable": "$.task_type", "StringEquals": "check_inventory", "Next": "InventoryAgent" }
      ],
      "Default": "RejectAndLog"
    },
    "PaymentAgent": {
      "Type": "Task",
      "Resource": "arn:aws:lambda:us-east-1:123456789012:function:payment-agent",
      "InputPath": "$.validated_params",
      "ResultPath": "$.payment_result",
      "TimeoutSeconds": 30,
      "Retry": [
        { "ErrorEquals": ["Lambda.TooManyRequestsException", "Lambda.ServiceException"],
          "IntervalSeconds": 2, "MaxAttempts": 3, "BackoffRate": 2.0 }
      ],
      "Catch": [
        { "ErrorEquals": ["States.ALL"], "ResultPath": "$.error", "Next": "RejectAndLog" }
      ],
      "Next": "Done"
    },
    "InventoryAgent": {
      "Type": "Task",
      "Resource": "arn:aws:lambda:us-east-1:123456789012:function:inventory-agent",
      "InputPath": "$.validated_params",
      "ResultPath": "$.inventory_result",
      "TimeoutSeconds": 30,
      "Retry": [
        { "ErrorEquals": ["Lambda.TooManyRequestsException", "Lambda.ServiceException"],
          "IntervalSeconds": 2, "MaxAttempts": 3, "BackoffRate": 2.0 }
      ],
      "Catch": [
        { "ErrorEquals": ["States.ALL"], "ResultPath": "$.error", "Next": "RejectAndLog" }
      ],
      "Next": "Done"
    },
    "RejectAndLog": {
      "Type": "Task",
      "Resource": "arn:aws:lambda:us-east-1:123456789012:function:log-rejected-message",
      "ResultPath": "$.log_result",
      "Next": "Rejected"
    },
    "Rejected": {
      "Type": "Fail",
      "Error": "MessageNotProcessed",
      "Cause": "Message failed validation, routing, or agent execution and was never forwarded."
    },
    "Done": {"Type": "Succeed"}
  }
}

Let me walk through what that buys you.

- One front door. Every message enters at ValidateRequest. Nothing reaches an agent until it has been validated and explicitly routed. There is no side entrance.
- Least-privilege payloads. InputPath: $.validated_params means the payment agent receives only the parameters it needs. It never sees the full envelope, the routing metadata, or the
nonce. If that Lambda gets popped, it has less to work with.
- Fail closed, and fail loud. Validation errors, unknown task_type values, and dead agents all funnel to RejectAndLog and then to a Fail state. This matters. In the naive version a
rejected message ends in Succeed, which means a malicious payload produces a green, successful execution and your CloudWatch alarm on failed executions never fires. Do not do that. A
rejection is a failure. Make it look like one.
- Retries for the honest failures. Agents throw transient errors. Lambda.TooManyRequestsException is not an attack, it is a Tuesday. The Retry block backs off and tries again before
giving up.
- A full audit trail. Every transition, with input and output, lands in the execution history. When something goes sideways at 3am you have a complete, replayable record of exactly what
happened. (Execution history is kept for 90 days, so if you need audit that outlives that, turn on CloudWatch Logs or ship it to S3. And remember the message contents sit in that
history, so mind what you log.)

Now about those timeouts, because the word "timeout" lies to you a little.

TimeoutSeconds: 30 does not kill your Lambda. Let me say that again, because I got this wrong myself the first time. When the timer fires, Step Functions stops waiting and marks the
task failed. The Lambda underneath keeps right on running, burning money and holding connections, until it finishes or hits its own function timeout. So set the Lambda's timeout too.
The state machine timeout protects the orchestrator. The function timeout protects your wallet and your downstream.

And here it comes. The part that actually makes this secure.

A state machine that routes through an orchestrator does not stop an agent from picking up the phone and calling another agent directly.

Read that again. The JSON above is necessary. It is not sufficient. If your payment agent's execution role has lambda:InvokeFunction on the inventory agent, then the moment someone
compromises the payment agent your beautiful orchestrator becomes optional. The attacker just calls the next agent directly and skips validation, routing, and logging. All of it.

The chokepoint is only a chokepoint if IAM makes it the only road.

So two rules, and they are not optional.

Rule 1: Only the orchestrator may invoke the agents. The state machine's execution role gets lambda:InvokeFunction, scoped to exactly these functions and nothing else.

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "InvokeOnlyTheseFunctions",
      "Effect": "Allow",
      "Action": "lambda:InvokeFunction",
      "Resource": [
        "arn:aws:lambda:us-east-1:123456789012:function:validate-agent-message",
        "arn:aws:lambda:us-east-1:123456789012:function:payment-agent",
        "arn:aws:lambda:us-east-1:123456789012:function:inventory-agent",
        "arn:aws:lambda:us-east-1:123456789012:function:log-rejected-message"
      ]
    }
  ]
}

Rule 2: The agents may not invoke each other. At all. Look closely at a typical agent role. Notice what is not there.

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "PaymentAgentDataAccessOnly",
      "Effect": "Allow",
      "Action": ["dynamodb:GetItem", "dynamodb:PutItem"],
      "Resource": "arn:aws:dynamodb:us-east-1:123456789012:table/payments"
    }
  ]
}

No lambda:InvokeFunction. Anywhere. The payment agent can touch its payments table and that is the entire list. It has no ability, none, to call the inventory agent or any other
function. If it needs something done elsewhere, it returns a result to the orchestrator and the orchestrator decides what happens next. That is what "no direct agent-to-agent
communication" actually means, and you will notice it lives in IAM, not in the state machine.

Want to go one step further? Put a resource-based policy on each agent function that grants invoke permission only to the orchestrator role, then use an SCP to deny
lambda:InvokeFunction on those ARNs for every other principal in the account. Now even an over-privileged admin role cannot accidentally wire two agents together. Belt, suspenders, and
a second belt.

The state machine gives you the shape. IAM gives you the guarantee. Most people ship only the first one and call it secure.

### 3. SQS Message Signing for Integrity

If your agents communicate via queues, [SQS](https://aws.amazon.com/sqs/?trk=d76afd77-bb62-46ac-b0a3-9dbf5ecde253) with message attributes and HMAC signing prevents tampering:

```python
import hashlib
import hmac
import json
import time
import uuid
import boto3

sqs = boto3.client("sqs")
secrets = boto3.client("secretsmanager")

# Pull the shared signing key from Secrets Manager, never from source control.
SECRET_KEY = secrets.get_secret_value(SecretId="agent-hmac-key")["SecretString"].encode()

MAX_AGE_SECONDS = 60  # freshness window; bounds the replay surface


class IntegrityError(Exception):
    """Raised when a message fails signature, freshness, or replay checks."""


class InMemoryNonceStore:
    """Example only. In production use a SHARED store (DynamoDB with a TTL
    attribute, or Redis). An in-memory dict does NOT stop replays across
    concurrent Lambda instances, because each instance has its own copy."""
    def __init__(self):
        self._seen = {}

    def exists(self, nonce: str) -> bool:
        now = time.time()
        self._seen = {n: exp for n, exp in self._seen.items() if exp > now}
        return nonce in self._seen

    def add(self, nonce: str, ttl: int):
        self._seen[nonce] = time.time() + ttl


def _signing_bytes(source_agent: str, nonce: str, timestamp: str, body: str) -> bytes:
    """Canonical, unambiguous representation of everything we sign. This covers
    the body AND the security-relevant attributes, so none of them can be
    tampered with independently of the signature."""
    return json.dumps(
        {"source_agent": source_agent, "nonce": nonce, "timestamp": timestamp, "body": body},
        sort_keys=True,
        separators=(",", ":"),
    ).encode()


def send_signed_message(queue_url: str, message: dict, source_agent: str):
    """Sign the body together with the attributes, then send."""
    body = json.dumps(message, sort_keys=True, separators=(",", ":"))
    nonce = str(uuid.uuid4())
    timestamp = str(int(time.time()))
    signature = hmac.new(
        SECRET_KEY, _signing_bytes(source_agent, nonce, timestamp, body), hashlib.sha256
    ).hexdigest()

    sqs.send_message(
        QueueUrl=queue_url,
        MessageBody=body,
        MessageAttributes={
            "source_agent": {"DataType": "String", "StringValue": source_agent},
            "signature":    {"DataType": "String", "StringValue": signature},
            "nonce":        {"DataType": "String", "StringValue": nonce},
            "timestamp":    {"DataType": "String", "StringValue": timestamp},
        },
    )


def _attr(attrs: dict, name: str) -> str:
    """Fetch a required string attribute, or fail closed."""
    try:
        return attrs[name]["StringValue"]
    except (KeyError, TypeError):
        raise IntegrityError(f"missing required attribute: {name}")


def verify_message(msg: dict, nonce_store) -> dict:
    """Verify a single message. Returns the parsed body or raises IntegrityError.
    `nonce_store` must expose .exists(nonce) and .add(nonce, ttl)."""
    body = msg["Body"]
    attrs = msg.get("MessageAttributes") or {}

    source_agent = _attr(attrs, "source_agent")
    signature    = _attr(attrs, "signature")
    nonce        = _attr(attrs, "nonce")
    timestamp    = _attr(attrs, "timestamp")

    # 1) Integrity + binding: recompute over body AND attributes, constant-time compare.
    expected = hmac.new(
        SECRET_KEY, _signing_bytes(source_agent, nonce, timestamp, body), hashlib.sha256
    ).hexdigest()
    if not hmac.compare_digest(signature, expected):
        raise IntegrityError(f"signature check failed (claimed source: {source_agent})")

    # 2) Freshness: reject stale or future-dated messages (allow a little clock skew).
    try:
        age = int(time.time()) - int(timestamp)
    except ValueError:
        raise IntegrityError("unparseable timestamp")
    if age > MAX_AGE_SECONDS or age < -MAX_AGE_SECONDS:
        raise IntegrityError(f"stale/future-dated message (age {age}s)")

    # 3) Anti-replay: each nonce is accepted exactly once.
    if nonce_store.exists(nonce):
        raise IntegrityError(f"replay detected (nonce {nonce})")
    nonce_store.add(nonce, ttl=MAX_AGE_SECONDS * 2)

    return json.loads(body)


def receive_and_process(queue_url: str, nonce_store, handler):
    """Long-poll, verify each message, process + delete on success. On failure
    we alert and leave the message for SQS to redrive to a dead-letter queue
    (configure a redrive policy with maxReceiveCount on the source queue)."""
    response = sqs.receive_message(
        QueueUrl=queue_url,
        MessageAttributeNames=["All"],
        MaxNumberOfMessages=10,
        WaitTimeSeconds=20,
    )
    for msg in response.get("Messages", []):
        try:
            payload = verify_message(msg, nonce_store)
        except IntegrityError as e:
            print(f"[REJECTED] {e}")   # emit a CloudWatch metric / alarm here
            continue                    # do NOT delete -> redrive policy sends it to the DLQ
        handler(payload)
        sqs.delete_message(QueueUrl=queue_url, ReceiptHandle=msg["ReceiptHandle"])