Test script for delegation cross account access by switching IAM Role.

test_switch_role.py

import boto3
import random


class S3(object):
    def __init__(self, credentials=None):
        if credentials:
            self.resource = boto3.resource('s3',
                    aws_access_key_id=credentials['AccessKeyId'],
                    aws_secret_access_key=credentials['SecretAccessKey'],
                    aws_session_token=credentials['SessionToken'],
                )
        else:
            self.resource = boto3.resource('s3')

    def create_bucket(self, bucket_name):
        self.resource.Bucket(bucket_name).create()

    def delete_bucket(self, bucket_name):
        self.delete_all_object(bucket_name)
        self.resource.Bucket(bucket_name).delete()

    def show_all_buckets(self):
        for bucket in self.resource.buckets.all():
            print(bucket.name)

    def show_all_keys(self, bucket_name):
        bucket = self.resource.Bucket(bucket_name)
        for key in bucket.objects.all():
            print(key.key)

    def upload_object(self, bucket_name, object_name, file_path):
        self.resource.Object(bucket_name, object_name).put(Body=open(file_path, 'r'))

    def delete_object(self, bucket_name, key_name):
        obj = self.resource.Object(bucket_name, key_name)
        obj.delete()

    def delete_all_object(self, bucket_name):
        bucket = self.resource.Bucket(bucket_name)
        for key in bucket.objects.all():
            key.delete()


def test_s3(s3):
    bucket_name = 'mybucket_for_testing_%s' % random.randint(0, 1000000)
    s3.create_bucket(bucket_name)
    s3.show_all_buckets()
    print('---- show current contents ---')
    s3.show_all_keys(bucket_name)
    print('--- upload testfile ---')
    s3.upload_object(bucket_name, 'testfile', './testfile')
    s3.show_all_keys(bucket_name)
    print('--- delete testfile ---')
    s3.delete_object(bucket_name, 'testfile')
    s3.show_all_keys(bucket_name)
    print('--- delete bucket ---')
    s3.delete_bucket(bucket_name)
    s3.show_all_buckets()
    
# access to default AWS account that is specified by environmental variables
s3 = S3()
test_s3(s3)

# get other aws's role credentials
sts = boto3.client('sts')
role_arn = 'arn:aws:iam::000000000000:role/test_role'

assume_role = sts.assume_role(
                RoleArn=role_arn,
                RoleSessionName='assume_role_session'
            )

print('--- switch role ---')
other_s3 = S3(assume_role['Credentials'])
test_s3(other_s3)