How to refresh the boto3 credetials when python script is running indefinitely

Viewed 6812

I am trying to write a python script that uses watchdog to look for file creation and upload that to s3 using boto3. However, my boto3 credentials expire after every 12hrs, So I need to renew them. I am storing my boto3 credentials in ~/.aws/credentials. So right now I am trying to catch the S3UploadFailedError, renew the credentials, and write them to ~/.aws/credentials. But though the credentials are getting renewed and I am calling boto3.client('s3') again its throwing exception.

What am I doing wrong? Or how can I resolve it?

Below is the code snippet

try:
     s3 = boto3.client('s3')
     s3.upload_file(event.src_path,'bucket-name',event.src_path)

except boto3.exceptions.S3UploadFailedError as e:
     print(e)
     get_aws_credentials()
     s3 = boto3.client('s3')

3 Answers

I have found a good example to refresh the credentials within this link: https://pritul95.github.io/blogs/boto3/2020/08/01/refreshable-boto3-session/

but there this a little bug inside. Be careful about that.

Here is the corrected code:

from uuid import uuid4
from datetime import datetime
from time import time

from boto3 import Session
from botocore.credentials import RefreshableCredentials
from botocore.session import get_session


class RefreshableBotoSession:
    """
    Boto Helper class which lets us create refreshable session, so that we can cache the client or resource.

    Usage
    -----
    session = RefreshableBotoSession().refreshable_session()

    client = session.client("s3") # we now can cache this client object without worrying about expiring credentials
    """

    def __init__(
        self,
        region_name: str = None,
        profile_name: str = None,
        sts_arn: str = None,
        session_name: str = None,
        session_ttl: int = 3000
    ):
        """
        Initialize `RefreshableBotoSession`

        Parameters
        ----------
        region_name : str (optional)
            Default region when creating new connection.

        profile_name : str (optional)
            The name of a profile to use.

        sts_arn : str (optional)
            The role arn to sts before creating session.

        session_name : str (optional)
            An identifier for the assumed role session. (required when `sts_arn` is given)

        session_ttl : int (optional)
            An integer number to set the TTL for each session. Beyond this session, it will renew the token.
            50 minutes by default which is before the default role expiration of 1 hour
        """

        self.region_name = region_name
        self.profile_name = profile_name
        self.sts_arn = sts_arn
        self.session_name = session_name or uuid4().hex
        self.session_ttl = session_ttl

    def __get_session_credentials(self):
        """
        Get session credentials
        """
        session = Session(region_name=self.region_name, profile_name=self.profile_name)

        # if sts_arn is given, get credential by assuming given role
        if self.sts_arn:
            sts_client = session.client(service_name="sts", region_name=self.region_name)
            response = sts_client.assume_role(
                RoleArn=self.sts_arn,
                RoleSessionName=self.session_name,
                DurationSeconds=self.session_ttl,
            ).get("Credentials")

            credentials = {
                "access_key": response.get("AccessKeyId"),
                "secret_key": response.get("SecretAccessKey"),
                "token": response.get("SessionToken"),
                "expiry_time": response.get("Expiration").isoformat(),
            }
        else:
            session_credentials = session.get_credentials().__dict__
            credentials = {
                "access_key": session_credentials.get("access_key"),
                "secret_key": session_credentials.get("secret_key"),
                "token": session_credentials.get("token"),
                "expiry_time": datetime.fromtimestamp(time() + self.session_ttl).isoformat(),
            }

        return credentials

    def refreshable_session(self) -> Session:
        """
        Get refreshable boto3 session.
        """
        # get refreshable credentials
        refreshable_credentials = RefreshableCredentials.create_from_metadata(
            metadata=self.__get_session_credentials(),
            refresh_using=self.__get_session_credentials,
            method="sts-assume-role",
        )

        # attach refreshable credentials current session
        session = get_session()
        session._credentials = refreshable_credentials
        session.set_config_variable("region", self.region_name)
        autorefresh_session = Session(botocore_session=session)

        return autorefresh_session

According to the documentation, the client looks in several locations for credentials and there are other options that are also more programmatic-friendly that you might want to consider instead of the .aws/credentials file.

Quoting the docs:

The order in which Boto3 searches for credentials is:

  • Passing credentials as parameters in the boto.client() method
  • Passing credentials as parameters when creating a Session object
  • Environment variables
  • Shared credential file (~/.aws/credentials)
  • AWS config file (~/.aws/config)
  • Assume Role provider

In your case, since you are already catching the exception and renewing the credentials, I would simply pass the new ones to a new instance of the client like so:

client = boto3.client(
    's3',
    aws_access_key_id=NEW_ACCESS_KEY,
    aws_secret_access_key=NEW_SECRET_KEY,
    aws_session_token=NEW_SESSION_TOKEN
)

If instead you are using these same credentials elsewhere in the code to create other clients, I'd consider setting them as environment variables:

import os

os.environ['AWS_ACCESS_KEY_ID'] = NEW_ACCESS_KEY
os.environ['AWS_SECRET_ACCESS_KEY'] = NEW_SECRET_KEY
os.environ['AWS_SESSION_TOKEN'] = NEW_SESSION_TOKEN

Again, quoting the docs:

The session key for your AWS account [...] is only needed when you are using temporary credentials.

Here is my implementation which only generates new credentials if existing credentials expire using a singleton design pattern

import boto3
from datetime import datetime
from dateutil.tz import tzutc
import os
import binascii


class AssumeRoleProd:
    __credentials = None

    def __init__(self):
        assert True==False

    @staticmethod
    def __setCredentials():
        print("\n\n ======= GENERATING NEW SESSION TOKEN ======= \n\n")
        # create an STS client object that represents a live connection to the
        # STS service
        sts_client = boto3.client('sts')

        # Call the assume_role method of the STSConnection object and pass the role
        # ARN and a role session name.
        assumed_role_object = sts_client.assume_role(
            RoleArn=your_role_here,
            RoleSessionName=f"AssumeRoleSession{binascii.b2a_hex(os.urandom(15)).decode('UTF-8')}"
        )

        # From the response that contains the assumed role, get the temporary
        # credentials that can be used to make subsequent API calls
        AssumeRoleProd.__credentials = assumed_role_object['Credentials']

    @staticmethod
    def getTempCredentials():
        credsExpired = False

        # Return object for the first time 
        if AssumeRoleProd.__credentials is None:
            AssumeRoleProd.__setCredentials()
            credsExpired = True

        # Generate if only 5 minutes are left for expiry. You may setup for entire 60 minutes by catching botocore ClientException
        elif (AssumeRoleProd.__credentials['Expiration']-datetime.now(tzutc())).seconds//60<=5: 
            AssumeRoleProd.__setCredentials()
            credsExpired = True

        return AssumeRoleProd.__credentials

And then I am using singleton design pattern for client as well which would generate a new client only if new session is generated. You can add region as well if required.

class lambdaClient:
    __prodClient = None

    def __init__(self):
        assert True==False

    @staticmethod
    def __initProdClient():
        credsExpired, credentials =  AssumeRoleProd.getTempCredentials()
        if lambdaClient.__prodClient is None or credsExpired:
            lambdaClient.__prodClient = boto3.client('lambda',
                                                       aws_access_key_id=credentials['AccessKeyId'],
                                                       aws_secret_access_key=credentials['SecretAccessKey'],
                                                       aws_session_token=credentials['SessionToken'])
        return lambdaClient.__prodClient


    @staticmethod
    def getProdClient():
        return lambdaClient.__initProdClient()
Related