WIP: Add AWS_SESSION_TOKEN parameter on boto3.client

ridi · Apr 5, 2021 · d1e3cd4 · d1e3cd4
1 parent 152eee1
commit d1e3cd4
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/ridi/secret_keeper/config.py b/ridi/secret_keeper/config.py
@@ -5,27 +5,32 @@
 ENVNAME_AWS_ACCESS_KEY = "SECRETKEEPER_AWS_ACCESS_KEY"
 ENVNAME_AWS_SECRET_KEY = "SECRETKEEPER_AWS_SECRET_KEY"
 ENVNAME_AWS_REGION = "SECRETKEEPER_AWS_REGION"
+ENVNAME_AWS_SESSION_TOKEN = "SECRETKEEPER_AWS_SESSION_TOKEN"
 
 ConnectArgsType = namedtuple('ConnectArgsType', [
-    'aws_access_key', 'aws_secret_key', 'aws_region'
+    'aws_access_key', 'aws_secret_key', 'aws_region', 'aws_session_token'
 ])
 
 AWS_CONNECT_ARGS = None
 
 
-def configure(aws_access_key=None, aws_secret_key=None, aws_region=None):
+def configure(aws_access_key=None, aws_secret_key=None, aws_region=None, aws_session_token=None):
     global AWS_CONNECT_ARGS
     AWS_CONNECT_ARGS = ConnectArgsType(
         aws_access_key=aws_access_key,
         aws_secret_key=aws_secret_key,
         aws_region=aws_region,
+        aws_session_token=aws_session_token,
     )
 
 
 configure(
     aws_access_key=os.environ.get(ENVNAME_AWS_ACCESS_KEY),
     aws_secret_key=os.environ.get(ENVNAME_AWS_SECRET_KEY),
     aws_region=os.environ.get(ENVNAME_AWS_REGION),
+    aws_session_token=os.environ.get(ENVNAME_AWS_SESSION_TOKEN),
+
 )
 
-__all__ = ["configure", "ENVNAME_AWS_ACCESS_KEY", "ENVNAME_AWS_SECRET_KEY", "ENVNAME_AWS_REGION"]
+__all__ = ["configure", "ENVNAME_AWS_ACCESS_KEY", "ENVNAME_AWS_SECRET_KEY", "ENVNAME_AWS_REGION",
+           "ENVNAME_AWS_SESSION_TOKEN"]
diff --git a/ridi/secret_keeper/connect.py b/ridi/secret_keeper/connect.py
@@ -9,6 +9,7 @@ def _get_client():
         region_name=AWS_CONNECT_ARGS.aws_region,
         aws_access_key_id=AWS_CONNECT_ARGS.aws_access_key,
         aws_secret_access_key=AWS_CONNECT_ARGS.aws_secret_key,
+        aws_session_token=AWS_CONNECT_ARGS.aws_session_token,
     )
     return client