First-class support for AssumeRole in sessions

[copumpkin]

This is similar to my request for boto/boto#3381, but for the botocore credentials/session system.

When I first looked for this, I got my hopes up because I saw the AssumeRoleProvider in credentials.py, but then it turned out to be fairly awkward to use programmatically with dynamically specified role metadata, as it seemed to assume fairly deeply that you wanted to use it the way aws-cli does, via a config file and static credentials.

What I'd really like to see is composable/fully programmatic solution to this. It would probably use much of the same logic that's already in AssumeRoleProvider, except with fewer assumptions about where the AssumeRole metadata information is coming from, and an API that makes it easy to create new assumed sessions from existing ones.

For example, here's an API I might enjoy using:

# Gives me some default session using default credentials that have power 
# to AssumeRole into other accounts
session = botocore.session.get_session() 

session1 = session.assume_role_session('arn:aws:iam::1234567890:role/JumpRole',
    role_session_name='hostile-takeover')
session2 = session.assume_role_session('arn:aws:iam::1111111111:role/SomeOtherJumpRole',
    role_session_name='hostile-takeover')

# arn:aws:iam::1234567890:role/JumpRole has the power to itself assume a role
session3 = session1.assume_role_session('arn:aws:iam::222222222:role/JumpRole',
    role_session_name='hostile-takeover')

# Now we wait for a few hours
time.sleep(3600 * 5)

client3 = session3.create_client('ec2', region_name='us-west-2')

# This should work, and transparently refresh credentials as needed up the
# stack (in this case, two credential refreshes would be needed since we're
# two AssumeRoles deep)
print client3.describe_instances()

I don't really care much about the API specifics, but I do want the entire AssumeRole session information to be (at least optionally) programmatic, and not implicitly loaded from some config file. The current AssumeRoleProvider also expects a source_profile which makes it hard to stack these things as I show above. Ideally, this would also work nicely with other AssumeRole variants, but that's far less pressing for me.

cc @jamesls who I think wrote (or at least ported) AssumeRoleProvider in botocore.

[copumpkin]

I have a simple working prototype of my proposal above that I cobbled together by ripping much of AssumeRoleProvider apart to do the bare minimum. I can try to clean it up and make a proper PR if you think that at a high level this is suitable functionality for botocore.

[mtdowling]

Thanks for the suggestion! I think you make a really good point that the usability and composibility of the AssumeRoleProvider could be improved so that it could be used directly. Before adding methods to session that are specific to credential providers, I think the first thing we should do is address the AssumeRoleProvider and try to make it easier to use on its own. One thing to comes to mind is updating the AssumeRoleProvider to have the ability to pass in a function or credential provider rather than assume that credentials are static or provided from a config.

I have a simple working prototype of my proposal above that I cobbled together by ripping much of AssumeRoleProvider apart to do the bare minimum. I can try to clean it up and make a proper PR if you think that at a high level this is suitable functionality for botocore.

Yes, that would be awesome. If I understand correctly what you're saying, if you could, I think it would be best to have one PR that is just the AssumeRoleProvider cleanup, and another PR that adds methods to Session. I know the former is something we are already sold on adding, and the latter would be something we'll discuss on the team and on the PR.

I'll make this as a feature request so that this is tracked. Thanks!

[copumpkin]

@mtdowling thanks for the prompt reply! My main concern about my current "speculative" code is that I basically just trimmed down the current AssumeRoleProvider (in my local project) for the sake of making it do what I want. I don't have easy access to the code from here right now but should be able to post it in the next day or two.

A preview of what I did until I can get back to the code:

1. Grab all of the current AssumeRoleProvider code and copy it to my own project
2. Make it so its constructor takes all the standard AssumeRole parameters and a pre-existing session to use for STS (the "parent")
3. Hack out everything beyond the basic core functionality so that all it has left is an in-memory cache, the refresh-from-STS code, and refresh avoidance using the cache. Things like the custom cache support, profile, source_profile, and the config loader are all gone.
4. Make a subclass of Session that hacks my custom AssumeRoleProvider into the credential provider (this seemed unexpectedly painful given all the component injection machinery I found in there; I might have screwed something up)

With the above, I can basically do everything I describe in my snippet above, except that instead of calling session.assume_role_session(some_arn, some_session_name, ...) I call AssumeRoleSession(session, some_arn, some_session_name, ...) .

More concrete things to do:

1. Given your concerns around adding to the Session API (which I understand), it needs to be easier to inject a custom CredentialProvider (like mine, or an instance of a more standard one that you provide) into a Session. I'd probably need your input on what that might look like.
2. I don't have any real use cases of the existing AssumeRoleProvider I can see outside of here and the test for it. How bad is it if I change its external API and just make sure the two use sites still work? I'm inclined to say not too bad because of how hard it is to use outside of your codebase, but if there are other projects using it, then I'd need to provide a shim, which is more work.

How about I start it as a gist showing the simple changes I describe, you take a look at what I did to inject my custom AssumeRoleProvider into a Session object, and then we decide what to do about point 1 above. For point 2, I basically just need your go-ahead on the API change. I can handle the surrounding fixes.

[copumpkin]

@mtdowling here's the code I was talking about: https://gist.github.com/copumpkin/1f8231d959d62934cb18. It's currently structured to work outside of the botocore repo, but I'd obviously prefer to move it into your codebase.

You can see some of the nasty voodoo I had to do at the bottom to actually inject my custom provider into a session. If there's a better way, please let me know, but I couldn't see an obvious way to register the component from outside without accessing its private member variable _components.

[mtdowling]

Thanks for putting together a gist. I'll be looping in @jamesls to help take a look at this as well. He wrote the original provider and should be able to provide some good feedback on what requirements we would have with regards to an updated AssumeRoleProvider.

[copumpkin]

@mtdowling @jamesls thanks! To clarify how I would replicate today's behavior, I would probably take my new AssumeRoleProvider and layer it on top of a SharedCredentialProvider session with the profile pointing at source_profile. At the same time, that feels a bit heavyweight, so if you have a better suggestion I'm open to that!

[kapilt]

Fwiw, i had to deal with the same problem, i ended up just reusing and overriding refreshable provider. I just did it as a function that takes a session and returns a wrapped session, basically same signature as the method here. Works well for chaining instance role to sts role for long running operations, also considerably less code. https://gist.github.com/kapilt/ac8e222081f63ba64e93

[update] hmm.. looking at assume role provider, i hadn't seen this one before, it looks like its handling a few more use cases, mfa, and cached sessions credentials. usage does look a bit akward though as its cli oriented with config via profile config against the session (most common lookup via file).

[copumpkin]

@kapilt thanks! I was considering something like what you did, but wanted to minimize the diff from the current version of the code to simplify review. I'd probably simplify it to look closer to yours in the longer run though.

@jamesls have you had a chance to take a look at this? I'd like to fix up my code a bit but would prefer to get some input before diving in.

[copumpkin]

@mtdowling is there anything I should be doing to get feedback? Should I just assume it's safe to do what I was proposing and submit a PR? I'd rather not do throwaway work if it's not going to get merged, which is why I hesitate to do that without prior feedback.

[mtdowling]

@copumpkin Sorry for the long delay. I'm not the best person to evaluate this PR right now, so I'll defer to @jamesls.

[pmdumuid]

Coming up to a year.. Any updates on this @mtdowling / @jamesls ?

[eedwardsdisco]

+1

[carlsverre]

👍 This would be super useful for writing CI/CD runners that assume various roles in order to deploy things.

[patrobinson]

FYI I got around this by using python cachetools library to memoize a function call with a TTL set to less than the credential timeout.

[ghost]

When assuming a role, it'd be great to have the limit for the token be set to a time greater than one hour. It's problematic for long running processes and it causes issues and headaches that as a developer you implicitly would appreciate if it did out of the box.

[jorisd]

@BardiaAfshin this 1hour token expiration is shit, but it's a AWS limitation AFAIK...http://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRole.html#API_AssumeRole_RequestParameters

[max-allan-surevine]

@jorisd Which 1hour limit ?

The duration, in seconds, of the role session. The value can range from 900 seconds (15 minutes) up to the maximum session duration setting for the role. This setting can have a value from 1 hour to 12 hours.

[patrobinson]

@max-allan-surevine that change was only recent (this year). Before that the 1 hour limit was a hard limit you could not change.

[jaymccon]

+1 on this, having to setup a config file to do this is not ideal in our use-case, would be great to be able to session = botocore.session.get_session(role_arn="arn:aws:iam::1111111111:role/MyRole") and get the same handling of credential rotation as with the AssumeRoleProvider

[jstewmon]

FWIW, the code required to assume role has simplified over time. Here's an example from one of my apps:

from botocore.credentials import (
    AssumeRoleCredentialFetcher,
    CredentialResolver,
    DeferredRefreshableCredentials,
    JSONFileCache
)
from botocore.session import Session


class AssumeRoleProvider(object):
    METHOD = 'assume-role'

    def __init__(self, fetcher):
        self._fetcher = fetcher

    def load(self):
        return DeferredRefreshableCredentials(
            self._fetcher.fetch_credentials,
            self.METHOD
        )


def assume_role(session: Session,
                role_arn: str,
                duration: int = 3600,
                session_name: str = None,
                serial_number: str = None) -> Session:
    fetcher = AssumeRoleCredentialFetcher(
        session.create_client,
        session.get_credentials(),
        role_arn,
        extra_args={
            'DurationSeconds': duration,
            'RoleSessionName': session_name,
            'SerialNumber': serial_number
        },
        cache=JSONFileCache()
    )
    role_session = Session()
    role_session.register_component(
        'credential_provider',
        CredentialResolver([AssumeRoleProvider(fetcher)])
    )
    return role_session

[vmksuw]

+1 for this.

[lightningboltemoji]

By my count, 6 out of 8 of the official AWS SDKs support programatic, automatically-refreshing "assume role" credentials - the only ones that don't are Python and PHP. I know it's already linked by GitHub, but I just wanted to quickly mention that I compiled examples/documentation for how this feature works in (almost) all other AWS SDKs: boto/boto3#3143

Just wanted to make sure that kind of "bigger picture" context got considered as well.

[mrmilosz]

Here's my take on @jstewmon's implementation:

import boto3
import botocore.credentials
import botocore.session

class AssumeRoleCredentialProvider(botocore.credentials.CredentialProvider):
    METHOD = "assume-role"
    CANONICAL_NAME = "custom-assume-role-provider"

    def __init__(self, session: boto3.Session, **assume_role_parameters):
        super().__init__()
        self._fetcher = botocore.credentials.AssumeRoleCredentialFetcher(
            client_creator=session.client,
            source_credentials=session.get_credentials(),
            role_arn=assume_role_parameters["RoleArn"],
            extra_args={
                key: value
                for key, value in assume_role_parameters.items()
                if key != "RoleArn"
            }
        )

    def load(self):
        return botocore.credentials.DeferredRefreshableCredentials(
            refresh_using=self._fetcher.fetch_credentials,
            method=self.METHOD
        )


def assume_role(session: boto3.Session, **assume_role_parameters) -> boto3.Session:
    botocore_session = botocore.session.Session()
    botocore_session.register_component(
        "credential_provider",
        botocore.credentials.CredentialResolver(
            [AssumeRoleCredentialProvider(session, **assume_role_parameters)]
        )
    )
    return boto3.Session(botocore_session=botocore_session)

It operates on a Boto3 session instead of a Botocore session and makes the parameters of assume_role equivalent to the parameters of an "assume-role" invocation on a regular Boto3 STS client. Usage example:

session = assume_role(
    session,
    RoleArn="arn:some_role_arn",
    RoleSessionName="SomeRoleSessionName",
    DurationSeconds=3600
)

session = assume_role(
    session,
    RoleArn="arn:some_other_role_arn",
    RoleSessionName="SomeOtherRoleSessionName",
    DurationSeconds=3600
)

Suppose you wanted to use this to call an AWS API Gateway API using an HTTP client:

auth = requests_aws4auth.AWS4Auth(
    region=session.region_name,
    service="execute-api",
    refreshable_credentials=session.get_credentials()
)

response = requests.post("https://some_url", auth=auth)

I think this could be easily ported to the Boto3 STS resource. Usage of this interface might look like this:

session = boto3.Session()
session = session.resource('sts').AssumeRoleSession(
    RoleArn="arn:some_role_arn",
    RoleSessionName="SomeRoleSessionName",
    DurationSeconds=3600
)

[benkehoe]

I have a comprehensive implementation with aws-assume-role-lib, including type annotations, parameter validation, allowing nicer parameter types (policy as JSON, duration as timedelta, etc). It can also monkeypatch boto3 to add an assume_role() method to boto3.Session as well as a top-level boto3.assume_role() that operates on the default session like boto3.client().
https://github.com/benkehoe/aws-assume-role-lib

It should be a candidate for including in boto3, but I've had a pull request for adding the credential provider open and gotten no response in nearly two years. #2096

[Azmisov]

Here is a version that will work for aioboto3 (and can be modified slightly for plain aiobotocore):

import aioboto3, asyncio
from aiobotocore.credentials import (
	AioAssumeRoleCredentialFetcher,
	AioDeferredRefreshableCredentials,
	AioCredentialResolver
)
from aiobotocore.session import get_session
from botocore.credentials import _local_now

class AioAssumeRoleCredentialProvider:
	METHOD = "assume-role"
	CANONICAL_NAME = "aio-assume-role-provider"
	
	def __init__(self, client_creator, source_credentials, RoleArn, **kwargs):
		self._fetcher = AioAssumeRoleCredentialFetcher(
			client_creator=client_creator,
			source_credentials=source_credentials,
			role_arn=RoleArn,
			extra_args=kwargs
		)

	async def load(self):
		return AioDeferredRefreshableCredentials(
			refresh_using=self._fetcher.fetch_credentials,
			method=self.METHOD,
			time_fetcher=_local_now
		)

async def assume_role (session: aioboto3.Session, **kwargs) -> aioboto3.Session:
	""" Assume role for an aioboto3 session, with autocredential refresh;
		See https://github.com/boto/botocore/issues/761
	"""
	parent = session._session
	# create autorefresh credential provider
	client_creator = parent.create_client
	source_credentials = await parent.get_credentials ()
	provider = AioAssumeRoleCredentialProvider (client_creator, source_credentials, **kwargs)
	# create child session which will assume the role
	child = get_session()
	child.register_component("credential_provider", AioCredentialResolver([provider]))
	# convert to aioboto3
	return aioboto3.Session(botocore_session=child)