Boto sessions and AWS multi-account

Generally when I’m writing an automation script for AWS resources, the action is isolated to the one account. Recently, I worked on a script that manipulated resources across multiple accounts. It’s good practice and a common pattern to host separate environments and resources in different accounts, unifying them then by creating a third. Users are assigned limited permissions in the third account, they can then take the step of assuming an IAM role to access the secure environments through a trust relationship with it. The benefits of this decoupling are:

  • it limits blast radius: if something fails in one account the other accounts are protected in their isolation;
  • you can host sensitive resources in one account and share them over in certain situations, limiting the risk and visibility you expose them to;
  • you can issue top-down permissions from other accounts from another more authoritative account. This was brought in with AWS Organisations becoming generally available;
  • comprehensibility: it’s easier to understand what’s going on in an environment when resources are not grouped together in one place.

With this in mind I wrote a quick Python module that can be imported when you need to jump from an authentication account into a more secure account through the assume_role method. It also handles the case that you are using multi-factor authentication, which is pretty important when running scripts against your infrastructure.

It can be used in the following manner:

account_session_a = Sts(
    role_arn='arn:aws:iam::<accnoa>:role/<rolename>',
    temporary_credentials_path='/tmp/creds_a.json',
    mfa_arn='arn:aws:iam::<accnoa>:mfa/<username>'
)
account_session_b = Sts(
    role_arn='arn:aws:iam::<accnob>:role/<rolename>',
    temporary_credentials_path='/tmp/creds_b.json',
    mfa_arn='arn:aws:iam::<accnob>:mfa/<username>'
)

account_session_a.client('sts').get_caller_identity().get('Account')
>>> account_id_a
account_session_b.client('sts').get_caller_identity().get('Account')
>>> account_id_b

The good bit about this code is if you are debugging script behaviour and you need to run it multiple times this code will test the credentials saved to disk from an earlier run, only if they have expired will they reprompt you to re-enter your MFA token (testing scripts should only ever be done against development environments, obviously). In the code I’ve set the expiry to 15 minutes (900 seconds), however this can be extended up to an hour (3600 seconds).

This also assumes you have previously set credentials for the authentication account you first operate in through ~/.aws/credentials.

# -*- coding: utf-8 -*-

import json
import boto3


from botocore.exceptions import ClientError


class Sts(object):
    """
    Sts: Object to manage the persistence of authentication over multiple
        runs of an automation script. When testing a script this will
        save having to input an MFA token multiple times when using
        an account that requires it.
    """

    def __init__(self, role_arn, temporary_credentials_path, mfa_arn):
        self.temp_creds_path = temporary_credentials_path
        self.mfa_arn = mfa_arn
        self.role_arn = role_arn

    def get_temporary_session(self):
        """
        get_temporary_session: checks the temporary credentials stored
            on disk, if they fail to authenticate re-attempt to assume
            the role. The credentials requested last 15 minutes. For
            debugging purposes these can be persisted for up to an hour.
        """

        try:
            with open(self.temp_creds_path, 'r') as tmp_creds:
                credentials = json.loads(tmp_creds.read())
                client = boto3.client(
                    "sts",
                    aws_access_key_id=credentials['AccessKeyId'],
                    aws_secret_access_key=credentials['SecretAccessKey'],
                    aws_session_token=credentials['SessionToken']
                )
                _ = client.get_caller_identity()["Account"]
        except (IOError, ClientError):
            response = boto3.client('sts').assume_role(
                DurationSeconds=900,
                RoleArn=self.role_arn,
                RoleSessionName='multi-account automation script',
                SerialNumber=self.mfa_arn,
                TokenCode=raw_input('MFA_Token:')
            )
            credentials = response['Credentials']
            with open(self.temp_creds_path, 'w') as tmp_creds:
                tmp_creds.write(json.dumps({
                    'AccessKeyId': credentials['AccessKeyId'],
                    'SecretAccessKey': credentials['SecretAccessKey'],
                    'SessionToken': credentials['SessionToken']}))

        return boto3.Session(
            aws_access_key_id=credentials['AccessKeyId'],
            aws_secret_access_key=credentials['SecretAccessKey'],
            aws_session_token=credentials['SessionToken'],
        )

Give this a go, if you have any issues or questions shoot me a message in the comments below, thanks!

comments powered by Disqus