assume_role.py - AWS Code Sample

assume_role.py

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: Apache-2.0 """ Purpose Shows how to assume a role using AWS Security Token Service (AWS STS) credentials. """ import json import time import sys import boto3 from botocore.exceptions import ClientError def progress_bar(seconds): """Shows a simple progress bar in the command window.""" for _ in range(seconds): time.sleep(1) print('.', end='') sys.stdout.flush() print() def unique_name(base_name): return f'demo-assume-role-{base_name}-{time.time_ns()}' def setup(iam_resource): """ Creates a new user with no permissions. Creates an access key pair for the user. Creates a role with a policy that lets the user assume the role. Creates a policy that allows listing Amazon S3 buckets. Attaches the policy to the role. Creates an inline policy for the user that lets the user assume the role. For demonstration purposes, the user is created in the same account as the role, but in practice the user would likely be from another account. :param iam_resource: A Boto3 AWS Identity and Access Management (IAM) resource that has permissions to create users, roles, and policies in the account. :return: The newly created user, user key, and role. """ user = iam_resource.create_user(UserName=unique_name('user')) print(f"Created user {user.name}.") user_key = user.create_access_key_pair() print(f"Created access key pair for user.") print(f"Wait for user to be ready.", end='') progress_bar(10) role = iam_resource.create_role( RoleName=unique_name('role'), AssumeRolePolicyDocument=json.dumps({ 'Version': '2012-10-17', 'Statement': [ { 'Effect': 'Allow', 'Principal': {'AWS': user.arn}, 'Action': 'sts:AssumeRole' } ] }) ) print(f"Created role {role.name}.") policy = iam_resource.create_policy( PolicyName=unique_name('policy'), PolicyDocument=json.dumps({ 'Version': '2012-10-17', 'Statement': [ { 'Effect': 'Allow', 'Action': 's3:ListAllMyBuckets', 'Resource': 'arn:aws:s3:::*' } ] }) ) role.attach_policy(PolicyArn=policy.arn) print(f"Created policy {policy.policy_name} and attached it to the role.") user.create_policy( PolicyName=unique_name('user-policy'), PolicyDocument=json.dumps({ 'Version': '2012-10-17', 'Statement': [ { 'Effect': 'Allow', 'Action': 'sts:AssumeRole', 'Resource': role.arn } ] }) ) print(f"Created an inline policy for {user.name} that lets the user assume " f"the role.") print("Give AWS time to propagate these new resources and connections.", end='') progress_bar(10) return user, user_key, role def show_access_denied_without_role(user_key): """ Shows that listing buckets without first assuming the role is not allowed. :param user_key: The key of the user created during setup. This user does not have permission to list buckets in the account. """ print(f"Try to list buckets without first assuming the role.") s3_denied_resource = boto3.resource( 's3', aws_access_key_id=user_key.id, aws_secret_access_key=user_key.secret) try: for bucket in s3_denied_resource.buckets.all(): print(bucket.name) raise RuntimeError("Expected to get AccessDenied error when listing buckets!") except ClientError as error: if error.response['Error']['Code'] == 'AccessDenied': print("Attempt to list buckets with no permissions: AccessDenied.") else: raise def list_buckets_from_assumed_role(user_key, assume_role_arn, session_name): """ Assumes a role that grants permission to list the Amazon S3 buckets in the account. Uses the temporary credentials from the role to list the buckets that are owned by the assumed role's account. :param user_key: The access key of a user that has permission to assume the role. :param assume_role_arn: The Amazon Resource Name (ARN) of the role that grants access to list the other account's buckets. :param session_name: The name of the STS session. """ sts_client = boto3.client( 'sts', aws_access_key_id=user_key.id, aws_secret_access_key=user_key.secret) response = sts_client.assume_role( RoleArn=assume_role_arn, RoleSessionName=session_name) temp_credentials = response['Credentials'] print(f"Assumed role {assume_role_arn} and got temporary credentials.") # Create an S3 resource that can access the account with the temporary credentials. s3_resource = boto3.resource( 's3', aws_access_key_id=temp_credentials['AccessKeyId'], aws_secret_access_key=temp_credentials['SecretAccessKey'], aws_session_token=temp_credentials['SessionToken']) print(f"Listing buckets for the assumed role's account:") for bucket in s3_resource.buckets.all(): print(bucket.name) def teardown(user, role): """ Removes all resources created during setup. :param user: The demo user. :param role: The demo role. """ for attached in role.attached_policies.all(): policy_name = attached.policy_name role.detach_policy(PolicyArn=attached.arn) attached.delete() print(f"Detached and deleted {policy_name}.") role.delete() print(f"Deleted {role.name}.") for user_pol in user.policies.all(): user_pol.delete() print("Deleted inline user policy.") for key in user.access_keys.all(): key.delete() print("Deleted user's access key.") user.delete() print(f"Deleted {user.name}.") def usage_demo(): """Drives the demonstration.""" print('-'*88) print(f"Welcome to the AWS Security Token Service assume role demo.") print('-'*88) iam_resource = boto3.resource('iam') user, user_key, role = setup(iam_resource) print(f"Created {user.name} and {role.name}.") try: show_access_denied_without_role(user_key) list_buckets_from_assumed_role(user_key, role.arn, 'AssumeRoleDemoSession') finally: teardown(user, role) print("Thanks for watching!") if __name__ == '__main__': usage_demo()