Bedrock AgentCore Policyで自然言語からCedarポリシーを生成しエージェントのツールアクセスを制御する

action  == AgentCore::Action::"[ターゲット名]___[ツール名]"
resource == AgentCore::Gateway::"[GatewayのARN]"

mkdir agentcore-policy-quickstart && cd agentcore-policy-quickstart
python3 -m venv .venv && source .venv/bin/activate
pip install boto3 bedrock-agentcore-starter-toolkit requests

import boto3
from bedrock_agentcore_starter_toolkit.operations.gateway.client import GatewayClient
from bedrock_agentcore_starter_toolkit.operations.policy.client import PolicyClient
 
gw_client = GatewayClient(region_name="us-west-2")
policy_client = PolicyClient(region_name="us-west-2")
control = boto3.client("bedrock-agentcore-control", region_name="us-west-2")
 
# 1. Policy Engine作成（名前はアンダースコアのみ許可、ハイフン不可）
engine = policy_client.create_or_get_policy_engine(
    name="my_policy_engine",
    description="Policy engine for tool access control",
)
engine_id = engine["policyEngineId"]
engine_arn = engine["policyEngineArn"]
 
# 2. MCP Gateway作成（名前はハイフンのみ許可、アンダースコア不可）
#    IAMロール・Cognito・Lambda関数が自動生成される
#    mode: "ENFORCE"（ポリシー強制）または "LOG_ONLY"（評価のみ、テスト用）
gateway = gw_client.create_mcp_gateway(
    name="my-gateway",
    policy_engine_config={
        "arn": engine_arn,
        "mode": "ENFORCE",
    },
)
gateway_id = gateway["gatewayId"]
gateway_arn = gateway["gatewayArn"]
gateway_url = gateway["gatewayUrl"]  # MCPエンドポイント
 
# 3. Lambda Targetの追加
#    target_payloadを省略するとデフォルトのget_weatherとget_timeツールが登録される
target = gw_client.create_mcp_gateway_target(
    gateway=gateway,
    name="my-target",
    target_type="lambda",
)

cedar_statement = f"""permit(
  principal,
  action == AgentCore::Action::"my-target___get_weather",
  resource == AgentCore::Gateway::"{gateway_arn}"
) when {{
  ((context.input).location) == "Tokyo"
}};"""
 
control.create_policy(
    name="allow_weather_tokyo",
    policyEngineId=engine_id,
    definition={"cedar": {"statement": cedar_statement}},
)

import requests, base64
 
cognito = boto3.client("cognito-idp", region_name="us-west-2")
 
# Gatewayの認証設定からCognito情報を取得
gw_detail = control.get_gateway(gatewayIdentifier=gateway_id)
auth_config = gw_detail["authorizerConfiguration"]["customJWTAuthorizer"]
discovery_url = auth_config["discoveryUrl"]
pool_id = discovery_url.split("/")[3]
client_id = auth_config["allowedClients"][0]
 
client_desc = cognito.describe_user_pool_client(UserPoolId=pool_id, ClientId=client_id)
client_secret = client_desc["UserPoolClient"]["ClientSecret"]
domain = cognito.describe_user_pool(UserPoolId=pool_id)["UserPool"]["Domain"]
token_endpoint = f"https://{domain}.auth.us-west-2.amazoncognito.com/oauth2/token"
 
# OAuthトークン取得（client_credentials grant）
auth_header = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode()
token = requests.post(token_endpoint, headers={
    "Content-Type": "application/x-www-form-urlencoded",
    "Authorization": f"Basic {auth_header}",
}, data={
    "grant_type": "client_credentials",
    "scope": f"{gateway['name']}/invoke",
}).json()["access_token"]

headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
 
def call_tool(name, arguments):
    return requests.post(gateway_url, headers=headers, json={
        "jsonrpc": "2.0",
        "method": "tools/call",
        "params": {"name": name, "arguments": arguments},
        "id": "1",
    })

{
  "tools": [
    {
      "name": "my-target___get_weather",
      "description": "Get weather for a location"
    }
  ]
}

{
  "result": {
    "isError": false,
    "content": [
      {
        "type": "text",
        "text": "{\"statusCode\":200,\"body\":\"{\\\"location\\\": \\\"Tokyo\\\", \\\"temperature\\\": \\\"72\\\\u00b0F\\\", \\\"conditions\\\": \\\"Sunny\\\"}\"}"
      }
    ]
  }
}

{
  "error": {
    "code": -32002,
    "message": "Tool Execution Denied: Tool call not allowed due to policy enforcement [No policy applies to the request (denied by default).]"
  }
}

gen = control.start_policy_generation(
    policyEngineId=engine_id,
    name="nl_generation_test",
    resource={"arn": gateway_arn},
    content={"rawText": "Allow calling the get_time tool only when the timezone is UTC"},
)
# ステータスがGENERATEDになるまでget_policy_generationでポーリング

permit(
  principal,
  action == AgentCore::Action::"my-target___get_time",
  resource == AgentCore::Gateway::"arn:aws:bedrock-agentcore:us-west-2:..."
) when {
  ((context.input).timezone) == "UTC"
};

policy_client.cleanup_policy_engine(engine_id)
gw_client.cleanup_gateway(gateway_id)