Manage Expiration of AWS Boto3 Assume Role Dynamically

Hi Folks ,

When automating tasks or performing CRUD operations with AWS services, it’s common to assume a role using AWS STS before creating service clients like EC2, S3, or others. The AssumeRole operation provides temporary credentials that expire after a specified duration, which can interrupt your automation tasks if not handled properly.

To ensure seamless execution of AWS operations, we can implement a Python decorator that checks the expiration of these temporary credentials before each service request. If the credentials have expired, the decorator will automatically reinitialize the Assume Role process and refresh the credentials.

This approach ensures that your service operations always have valid credentials without the need to manually recreate the service clients each time the token expires.

Here’s an example of how this can be achieved using Python, allowing you to dynamically handle the reauthentication process for any AWS service client:

import boto3
import datetime
from functools import wraps
from botocore.exceptions import ClientError

class STSAssumeRole:
    def __init__(self, role_arn, session_name, duration_seconds=3600):
        self.role_arn = role_arn
        self.session_name = session_name
        self.duration_seconds = duration_seconds
        self.credentials = None
        self.expiration = None
        self.assume_role()

    def assume_role(self):
        """Assume the specified IAM role and store the temporary credentials."""
        sts_client = boto3.client('sts')
        try:
            response = sts_client.assume_role(
                RoleArn=self.role_arn,
                RoleSessionName=self.session_name,
                DurationSeconds=self.duration_seconds
            )
            self.credentials = response['Credentials']
            self.expiration = self.credentials['Expiration']
            print(f"Assumed role successfully, credentials expire at {self.expiration}")
        except ClientError as e:
            print(f"Error assuming role: {e}")
            raise e

    def check_token_expiry(self):
        """Check if the token has expired and re-assume the role if necessary."""
        if self.expiration <= datetime.datetime.utcnow():
            print("Token expired, re-assuming role...")
            self.assume_role()

    def get_client(self, service_name):
        """Dynamically create and return a client for the specified AWS service."""
        self.check_token_expiry()
        return boto3.client(
            service_name,
            aws_access_key_id=self.credentials['AccessKeyId'],
            aws_secret_access_key=self.credentials['SecretAccessKey'],
            aws_session_token=self.credentials['SessionToken']
        )

# Decorator to handle token expiration before any AWS operation
def sts_role_checker(func):
    @wraps(func)
    def wrapper(instance, *args, **kwargs):
        instance.check_token_expiry()  # Check token before executing the function
        return func(instance, *args, **kwargs)
    return wrapper

# Example usage of the class with dynamic client creation and decorator
sts_assume_role = STSAssumeRole('arn:aws:iam::123456789012:role/MyRole', 'MySessionName')

@sts_role_checker
def list_ec2_instances(instance):
    try:
        ec2_client = instance.get_client('ec2')  # Dynamically create EC2 client
        response = ec2_client.describe_instances()
        return response
    except ClientError as e:
        print(f"Error listing EC2 instances: {e}")
        raise e

@sts_role_checker
def list_s3_buckets(instance):
    try:
        s3_client = instance.get_client('s3')  # Dynamically create S3 client
        response = s3_client.list_buckets()
        return response
    except ClientError as e:
        print(f"Error listing S3 buckets: {e}")
        raise e

# Calling the functions
instances = list_ec2_instances(sts_assume_role)
print(instances)

buckets = list_s3_buckets(sts_assume_role)
print(buckets)

#