gen_sts_creds.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.8"
# dependencies = [
#     "boto3",
#     "botocore",
# ]
# ///

"""
Standalone script to generate AWS STS temporary credentials.
Prints environment variable assignments to STDOUT for copy/paste.
"""

import sys
import boto3
import botocore.exceptions


def get_upload_credentials(role_session_name="upload-session", duration_seconds=3600):
    """
    Obtain temporary STS credentials by assuming the heapdump_upload_role.

    Args:
        role_session_name (str): Identifier for the session
        duration_seconds (int): Duration of the session in seconds (max 3600 as per role config)

    Returns:
        dict: Dictionary containing temporary credentials:
            - access_key: AWS access key ID
            - secret_key: AWS secret access key
            - session_token: AWS session token
            - expiration: Expiration timestamp of the credentials

    Raises:
        botocore.exceptions.ClientError: If role assumption fails
    """
    try:
        # Create an STS client
        session = boto3.Session()
        sts_client = session.client("sts")

        # Role ARN for the upload role - adjust account ID if needed
        role_arn = "arn:aws:iam::xxxx:role/heapdump-upload-role"

        # Assume the role
        print(f"Assuming role: {role_arn}", file=sys.stderr)
        response = sts_client.assume_role(
            RoleArn=role_arn,
            RoleSessionName=role_session_name,
            DurationSeconds=duration_seconds,
        )

        # Extract credentials from response
        credentials = response["Credentials"]

        print(f"Successfully obtained temporary credentials, valid until: {credentials['Expiration']}", file=sys.stderr)

        # Return credentials in a dictionary
        return {
            "access_key": credentials["AccessKeyId"],
            "secret_key": credentials["SecretAccessKey"],
            "session_token": credentials["SessionToken"],
            "expiration": credentials["Expiration"],
        }

    except botocore.exceptions.ClientError as e:
        print(f"Failed to assume role: {e}", file=sys.stderr)
        # If MFA is required but not provided
        if "MultiFactorAuthentication" in str(e):
            print("MFA is required. Please provide an MFA token.", file=sys.stderr)
        raise


def main():
    """Main function to generate and print STS credentials."""
    try:
        # Get temporary credentials
        upload_creds = get_upload_credentials()
    except botocore.exceptions.ClientError as e:
        print(f"Failed to get AWS STS upload credentials: {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"Unexpected error: {e}", file=sys.stderr)
        sys.exit(1)

    # Print export statements to STDOUT for copy/paste
    print(f"export AWS_ACCESS_KEY_ID={upload_creds['access_key']}")
    print(f"export AWS_SECRET_ACCESS_KEY={upload_creds['secret_key']}")
    print(f"export AWS_SESSION_TOKEN={upload_creds['session_token']}")


if __name__ == "__main__":
    main()