aws-assume-role-lib

The primary benefit of aws-assume-role-lib is that it handles the complexity of caching and refreshing assumed role credentials automatically.

When using aws_assume_role_lib.assume_role(), the returned session is valid for the life of the Lambda function instance because the library uses built-in boto3/botocore mechanisms to refresh credentials before they expire. This allows you to move the role assumption logic to the initialization phase (outside the handler) without worrying about credential expiration during long-lived function instances.

Key Features:

• Automatic Caching/Refreshing: No manual logic needed to check if credentials have expired.
• Automatic Session Naming: If RoleSessionName is absent, it is automatically generated.
• SourceIdentity Integration: If you use SourceIdentity, the library automatically sets the RoleSessionName to match it.

When called within a Lambda environment, generate_lambda_session_name() creates a session name that helps trace usage back to a specific function instance. It attempts to include the function name, version, and a unique identifier (from the CloudWatch log stream name).

The name is truncated to fit within the 64-character AWS limit. The format follows these priorities:

1. {function_name} (if name > 59 chars, it is truncated)
2. {function_name}.{identifier}
3. {function_name}.{function_version}.{identifier} (version is omitted if $LATEST)

If the identifier cannot be found, it falls back to a timestamp or a random string.

You can install the library using pip or by adding the source file directly to your project.

pip install --user aws-assume-role-lib

For most use cases, it is recommended to use the AWS SDK's built-in support for role assumption by configuring your ~/.aws/config file rather than using this library. This allows for automatic credential refreshing by the SDK.

To set up an assumed role profile, define a source_profile (an existing profile with credentials) or a credential_source (e.g., Environment, Ec2InstanceMetadata, or EcsContainer).

[profile profile-to-call-assume-role-with]
# Existing credentials (IAM User, SSO, etc.)

[profile my-assumed-role]
role_arn = arn:aws:iam::123456789012:role/MyRole
# optional: role_session_name = MyRoleSessionName
source_profile = profile-to-call-assume-role-with
# OR use credential_source = Environment | Ec2InstanceMetadata | EcsContainer

The demo uses aws-error-utils to catch common AWS service exceptions like AccessDenied more cleanly than checking botocore.exceptions.ClientError response codes manually.

Instead of inspecting the error code string (e.g., e.response["Error"]["Code"] == "AccessDenied"), you can catch the specific exception class provided by the library.

from aws_error_utils import errors

try:
    response = s3.get_object(
        Bucket=BUCKET_NAME,
        Key=KEY,
    )
    s3_result = response['Body'].read()
except errors.AccessDenied:
    s3_result = "Access denied!"
except Exception as e:
    s3_result = str(e)

To persist credentials on the file system, pass a JSONFileCache instance to the cache argument of the assume_role() function. The library will automatically create the directory specified in the constructor if it does not exist.

Alternatively, you can provide any dict-like object that supports __getitem__, __setitem__, and __contains__ to act as a cache.

assumed_role_session = assume_role(session, "arn:aws:iam::123456789012:role/MyRole", cache=JSONFileCache("path/to/dir"))

If you have the components of an ARN but need the exact string format, use the following utility functions:

• get_role_arn(): Constructs a standard IAM role ARN. If the role name includes a path, you can provide it within the role_name or via the path argument, but not both.
• get_assumed_role_session_arn(): Constructs an assumed role session ARN. Note that assumed role session ARNs do not include the role path; if a path is included in the role_name, it will be removed.

Both functions support specifying the partition (defaults to aws).

get_role_arn(
    account_id: Union[str, int],
    role_name:  str,
    path:       str             = "",
    partition:  str             = "aws",
)

get_assumed_role_session_arn(
    account_id:        Union[str, int],
    role_name:         str,
    role_session_name: str,
    partition:         str             = "aws",
)

The assume_role() function simplifies the process of assuming an IAM role by taking an existing boto3.Session and a RoleArn. It returns a new session that automatically handles credential refreshing when they expire. Unlike the standard STS assume_role API, you do not need to manually provide a RoleSessionName (it will be generated for you) or manually extract credentials to create a new session.

import boto3
from aws_assume_role_lib import assume_role

# Get a session
session = boto3.Session()

# Assume the session
assumed_role_session = assume_role(session, "arn:aws:iam::123456789012:role/MyRole")

# Use the assumed role session
print(assumed_role_session.client("sts").get_caller_identity()["Arn"])

You can use patch_boto3() to inject the assume_role functionality directly into the boto3 module. This adds boto3.assume_role(RoleArn, ...) (using the default session) and boto3.Session.assume_role(...) to your environment.

import boto3
import aws_assume_role_lib
aws_assume_role_lib.patch_boto3()

# Use the patched top-level function
assumed_role_session = boto3.assume_role("arn:aws:iam::123456789012:role/MyRole")

# Use the patched Session method
session = boto3.Session(profile_name="my-profile")
assumed_role_session = session.assume_role("arn:aws:iam::123456789012:role/MyRole")

To improve traceability in CloudTrail and CloudWatch, use aws_assume_role_lib.generate_lambda_session_name().

This function creates a session name that includes the function name, version, and function instance identifier. This allows you to correlate activity in CloudTrail logs directly back to specific CloudWatch Log streams.