Change credential cache directory

snigdhasjg · Nov 28, 2023 · d7bfe94 · d7bfe94
1 parent d049983
commit d7bfe94
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 8 deletions.
diff --git 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
@@ -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