[小ネタ] Amazon Bedrock AgentCore GatewayにIAM認証でクロスアカウントアクセスする

"""アカウントA（ツール提供側）にIAM認証のGatewayとLambdaターゲットを作成する"""

import json
import time
import io
import zipfile

import boto3

# === 設定 ===
REGION = "us-east-1"
PROFILE = "account-a"  # アカウントA（ツール提供側）のプロファイル名に変更してください
GATEWAY_NAME = "cross-account-test"

session = boto3.Session(profile_name=PROFILE, region_name=REGION)
iam = session.client("iam")
lambda_client = session.client("lambda", region_name=REGION)
agentcore = session.client("bedrock-agentcore-control", region_name=REGION)
account_id = session.client("sts").get_caller_identity()["Account"]

def create_gateway_service_role():
    """Gateway用のサービスロールを作成"""
    role_name = "AgentCoreGatewayServiceRole"
    trust_policy = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Principal": {"Service": "bedrock-agentcore.amazonaws.com"},
                "Action": "sts:AssumeRole",
            }
        ],
    }

    try:
        role = iam.create_role(
            RoleName=role_name,
            AssumeRolePolicyDocument=json.dumps(trust_policy),
        )
        print(f"Created role: {role['Role']['Arn']}")
    except iam.exceptions.EntityAlreadyExistsException:
        role = iam.get_role(RoleName=role_name)
        print(f"Role already exists: {role['Role']['Arn']}")

    return role["Role"]["Arn"]

def create_test_lambda():
    """テスト用のLambda関数を作成"""
    function_name = "cross-account-gateway-test"
    role_name = "CrossAccountTestLambdaRole"

    # Lambda実行ロール
    trust_policy = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Principal": {"Service": "lambda.amazonaws.com"},
                "Action": "sts:AssumeRole",
            }
        ],
    }

    try:
        role = iam.create_role(
            RoleName=role_name,
            AssumeRolePolicyDocument=json.dumps(trust_policy),
        )
        iam.attach_role_policy(
            RoleName=role_name,
            PolicyArn="arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole",
        )
        print(f"Created Lambda role: {role['Role']['Arn']}")
        print("Waiting for IAM propagation...")
        time.sleep(10)
    except iam.exceptions.EntityAlreadyExistsException:
        role = iam.get_role(RoleName=role_name)

    # Lambda関数コード
    code = '''
import json

def handler(event, context):
    return {
        "statusCode": 200,
        "body": json.dumps({"message": f"Hello from Gateway! Input: {event}"})
    }
'''
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
        zf.writestr("lambda_function.py", code)
    zip_buffer.seek(0)

    try:
        response = lambda_client.create_function(
            FunctionName=function_name,
            Runtime="python3.12",
            Handler="lambda_function.handler",
            Role=role["Role"]["Arn"],
            Code={"ZipFile": zip_buffer.read()},
        )
        print(f"Created Lambda: {response['FunctionArn']}")
    except lambda_client.exceptions.ResourceConflictException:
        response = lambda_client.get_function(FunctionName=function_name)
        response = response["Configuration"]
        print(f"Lambda already exists: {response['FunctionArn']}")

    return response["FunctionArn"]

def setup():
    print("=" * 60)
    print("Step 1: Gatewayサービスロールの作成")
    print("=" * 60)
    gateway_role_arn = create_gateway_service_role()

    print(f"\n{'=' * 60}")
    print("Step 2: テスト用Lambda関数の作成")
    print("=" * 60)
    lambda_arn = create_test_lambda()

    # GatewayロールにLambda Invoke権限を付与
    iam.put_role_policy(
        RoleName="AgentCoreGatewayServiceRole",
        PolicyName="LambdaInvoke",
        PolicyDocument=json.dumps({
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Action": "lambda:InvokeFunction",
                    "Resource": lambda_arn,
                }
            ],
        }),
    )
    print("Attached Lambda invoke policy to Gateway role")

    print(f"\n{'=' * 60}")
    print("Step 3: Gatewayの作成（IAM認証）")
    print("=" * 60)
    gateway = agentcore.create_gateway(
        name=GATEWAY_NAME,
        description="Cross-account access test gateway",
        protocolType="MCP",
        authorizerType="AWS_IAM",
        roleArn=gateway_role_arn,
    )
    gateway_id = gateway["gatewayId"]
    gateway_url = gateway["gatewayUrl"]
    gateway_arn = gateway["gatewayArn"]
    print(f"Gateway ID:  {gateway_id}")
    print(f"Gateway URL: {gateway_url}")
    print(f"Gateway ARN: {gateway_arn}")

    # READYになるまで待機
    print("Waiting for Gateway to become READY...")
    while True:
        status = agentcore.get_gateway(gatewayIdentifier=gateway_id)["status"]
        if status == "READY":
            break
        time.sleep(3)
    print("Gateway is READY!")

    print(f"\n{'=' * 60}")
    print("Step 4: Lambdaターゲットの追加")
    print("=" * 60)
    target = agentcore.create_gateway_target(
        gatewayIdentifier=gateway_id,
        name="echo-tools",
        targetConfiguration={
            "mcp": {
                "lambda": {
                    "lambdaArn": lambda_arn,
                    "toolSchema": {
                        "inlinePayload": [
                            {
                                "name": "echo_message",
                                "description": "Echo back a message for testing",
                                "inputSchema": {
                                    "type": "object",
                                    "properties": {
                                        "message": {
                                            "type": "string",
                                            "description": "The message to echo",
                                        }
                                    },
                                    "required": ["message"],
                                },
                            }
                        ]
                    },
                }
            }
        },
        credentialProviderConfigurations=[
            {"credentialProviderType": "GATEWAY_IAM_ROLE"}
        ],
    )
    target_id = target["targetId"]
    print(f"Target ID: {target_id}")

    # READYになるまで待機
    print("Waiting for Target to become READY...")
    while True:
        status = agentcore.get_gateway_target(
            gatewayIdentifier=gateway_id, targetId=target_id
        )["status"]
        if status == "READY":
            break
        time.sleep(3)
    print("Target is READY!")

    # 設定を保存
    config = {
        "gateway_id": gateway_id,
        "gateway_url": gateway_url,
        "gateway_arn": gateway_arn,
        "target_id": target_id,
        "lambda_arn": lambda_arn,
        "account_id": account_id,
        "region": REGION,
    }
    with open("gateway_config.json", "w") as f:
        json.dump(config, f, indent=2)

    print(f"\n{'=' * 60}")
    print("Setup complete!")
    print(f"Configuration saved to: gateway_config.json")
    print(f"Gateway URL: {gateway_url}")
    print("=" * 60)

if __name__ == "__main__":
    setup()

"""作成したリソースを削除する"""

import json
import boto3

REGION = "us-east-1"
PROFILE = "account-a"  # アカウントA（ツール提供側）

session = boto3.Session(profile_name=PROFILE, region_name=REGION)
agentcore = session.client("bedrock-agentcore-control", region_name=REGION)
lambda_client = session.client("lambda", region_name=REGION)
iam = session.client("iam")

with open("gateway_config.json") as f:
    config = json.load(f)

gateway_id = config["gateway_id"]
target_id = config["target_id"]
lambda_arn = config["lambda_arn"]

# リソースベースポリシーの削除
try:
    agentcore.delete_resource_policy(resourceArn=config["gateway_arn"])
    print("Deleted resource policy")
except Exception:
    pass

# ターゲットの削除
agentcore.delete_gateway_target(
    gatewayIdentifier=gateway_id, targetId=target_id
)
print(f"Deleted target: {target_id}")

# Gatewayの削除
agentcore.delete_gateway(gatewayIdentifier=gateway_id)
print(f"Deleted gateway: {gateway_id}")

# Lambda関数の削除
lambda_client.delete_function(FunctionName="cross-account-gateway-test")
print("Deleted Lambda function")

# IAMロールの削除
for role_name in ["CrossAccountTestLambdaRole", "AgentCoreGatewayServiceRole"]:
    try:
        # インラインポリシーの削除
        policies = iam.list_role_policies(RoleName=role_name)
        for policy_name in policies["PolicyNames"]:
            iam.delete_role_policy(RoleName=role_name, PolicyName=policy_name)
        # マネージドポリシーのデタッチ
        attached = iam.list_attached_role_policies(RoleName=role_name)
        for policy in attached["AttachedPolicies"]:
            iam.detach_role_policy(
                RoleName=role_name, PolicyArn=policy["PolicyArn"]
            )
        iam.delete_role(RoleName=role_name)
        print(f"Deleted role: {role_name}")
    except Exception as e:
        print(f"Error deleting {role_name}: {e}")

print("\nCleanup complete!")