Botocore: First-class support for AssumeRole in sessions

[copumpkin]

This is similar to my request for https://github.com/boto/boto/issues/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]

:+1: 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.

[BardiaAfshin]

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.