From d7bfe94b2089be29405244e768cd85c21680c6c5 Mon Sep 17 00:00:00 2001 From: Snigdhajyoti Ghosh Date: Tue, 28 Nov 2023 23:17:12 +0530 Subject: [PATCH] Change credential cache directory --- aws_fusion/aws/assume_role.py | 20 +++++++++++++------- aws_fusion/aws/session.py | 5 ++++- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/aws_fusion/aws/assume_role.py b/aws_fusion/aws/assume_role.py index b3f4bb8..a00816e 100644 --- a/aws_fusion/aws/assume_role.py +++ b/aws_fusion/aws/assume_role.py @@ -4,6 +4,7 @@ import datetime import logging import json +import os from botocore.exceptions import ClientError from botocore.utils import JSONFileCache @@ -12,7 +13,8 @@ class AssumeRoleWithSamlCache: - __jsonFileCache = JSONFileCache() + __CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'saml', 'cache')) + __jsonFileCache = JSONFileCache(__CACHE_DIR) def __init__(self, role) -> None: LOG.debug('Initialize AssumeRoleWithSamlCache') @@ -23,10 +25,13 @@ def __init__(self, role) -> None: def does_valid_token_cache_exists(self): if self.__cache_key in self.__jsonFileCache: response = self.__jsonFileCache[self.__cache_key] - expiration = datetime.datetime.strptime(response['Credentials']['Expiration'], '%Y-%m-%dT%H:%M:%S%Z') - current_utc_time = datetime.datetime.utcnow() + expiration = (datetime.datetime + .strptime(response['Credentials']['Expiration'], '%Y-%m-%dT%H:%M:%S%Z') + .replace(tzinfo=datetime.timezone.utc)) + current_utc_time = datetime.datetime.now(tz=datetime.timezone.utc) + if expiration - current_utc_time >= datetime.timedelta(minutes=1): - response['Credentials']['Expiration'] = expiration.replace(tzinfo=datetime.timezone.utc) + response['Credentials']['Expiration'] = expiration self.__response = response LOG.debug('Valid token exists. Can use cache') return True @@ -42,9 +47,10 @@ def credential_process(self): "Version": 1, "AccessKeyId": credentials['AccessKeyId'], "SecretAccessKey": credentials['SecretAccessKey'], - "SessionToken": credentials['SessionToken'] + "SessionToken": credentials['SessionToken'], + "Expiration": credentials['Expiration'].isoformat() }) - + def environment_variable(self): credentials = self.__response['Credentials'] LOG.debug(f'Giving credential as environment variable format') @@ -59,6 +65,7 @@ def environment_variable(self): def assume_role_with_saml(self, saml_response, roles, session_duration): LOG.debug(f'Started assuming role with SAML') + # Just need a dummy sts session client object to call assume role with saml client = boto3.Session(aws_access_key_id='dummy', aws_secret_access_key='dummy').client('sts') selected_role = self.__role try: @@ -82,4 +89,3 @@ def assume_role_with_saml(self, saml_response, roles, session_duration): self.__jsonFileCache[self.__cache_key] = response self.__response = response - diff --git a/aws_fusion/aws/session.py b/aws_fusion/aws/session.py index 92b7d29..170cf7d 100644 --- a/aws_fusion/aws/session.py +++ b/aws_fusion/aws/session.py @@ -1,10 +1,12 @@ import boto3 import logging +import os from botocore.utils import JSONFileCache LOG = logging.getLogger(__name__) + class TokenGenerationException(Exception): """Exception for credential not having token""" pass @@ -26,9 +28,10 @@ def credentials(profile_name, region_name): def __update_credential_provider_cache(session): """Setting up a custom cache implementation like aws cli""" + cache_dir = os.path.expanduser(os.path.join('~', '.aws', 'cli', 'cache')) cred_chain = session._session.get_component('credential_provider') - json_file_cache = JSONFileCache() + json_file_cache = JSONFileCache(cache_dir) def _update(provider_name): cred_chain.get_provider(provider_name).cache = json_file_cache