Refresh short-term Amazon Bedrock API keys for greater control and security

Once you are more familiar with Amazon Bedrock, we strongly recommend against using long-term Amazon Bedrock API keys. Instead, you should switch to using temporary security credentials or to short-term Amazon Bedrock API keys.

You can create a script with the help of the aws-bedrock-token-generator package to programmatically regenerate a new short-term key whenever your current one has expired. First, ensure that you've fulfilled the prerequisites at Generate a short-term Amazon Bedrock API key using the API.

You can then use the following Python script:

from datetime import datetime, timedelta
import os
import boto3
from botocore.credentials import Credentials
from aws_bedrock_token_generator import BedrockTokenGenerator

# Replace the following values as necessary
SESSION_DURATION = timedelta(hours=12)                                  # 12 hours is the maximum
EFFECTIVE_TOKEN_DURATION = min(SESSION_DURATION, timedelta(hours=12))   # The token can last no longer than 12 hours
ROLE_ARN = "arn:aws:iam::111122223333:role/TargetRole"        # Ensure that the identity you're authenticating with has permissions to assume this role
ROLE_SESSION_NAME = "your-session-name"
REGION = "us-east-1"

def get_session_from_assume():
    sts = boto3.client("sts")
    response = sts.assume_role(
        RoleArn=ROLE_ARN,
        RoleSessionName=ROLE_SESSION_NAME,
        DurationSeconds=int(SESSION_DURATION.total_seconds())
    )
    creds = response["Credentials"]
    return Credentials(
        access_key=creds["AccessKeyId"],
        secret_key=creds["SecretAccessKey"],
        token=creds["SessionToken"]
    )

# Step 1: Generate initial token and note timestamp
generator = BedrockTokenGenerator()
creds = get_session_from_assume()
token = generator.get_token(creds, region=REGION)
token_created_at = datetime.utcnow()

# Step 2: Later in your long lived process — before using the token
if datetime.utcnow() - token_created_at >= EFFECTIVE_TOKEN_DURATION:
    creds = get_session_from_assume()
    token = generator.get_token(creds, region=REGION)
    token_created_at = datetime.utcnow()
else:
    # Token is still valid
    pass
        
# Set the token as an environment variable so you can use it in API calls
os.environ['AWS_BEARER_TOKEN_BEDROCK'] = token