From af02258c693408e6791ad3149a40affec2707c94 Mon Sep 17 00:00:00 2001
From: Krzysztof Nazarewski <3494992+nazarewk@users.noreply.github.com>
Date: Thu, 21 Apr 2022 11:11:48 +0200
Subject: [PATCH] Configurable SSO directories

Allows configuring SSO directories explicitly in addition to more sane defaults (discovering them from AWS SDK credentials/config file locations)
---
 lib/aws_sso_lib/sso.py | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/lib/aws_sso_lib/sso.py b/lib/aws_sso_lib/sso.py
index 74de1c2..1cd3563 100644
--- a/lib/aws_sso_lib/sso.py
+++ b/lib/aws_sso_lib/sso.py
@@ -37,12 +37,22 @@
 from .exceptions import InvalidSSOConfigError, AuthDispatchError, AuthenticationNeededError
 from .browser import OpenBrowserHandler, non_interactive_auth_raiser
 
-SSO_TOKEN_DIR = os.path.expanduser(
-    os.path.join("~", ".aws", "sso", "cache")
+AWS_CREDENTIALS_DIR = (
+    os.path.dirname(os.environ.get("AWS_SHARED_CREDENTIALS_FILE", "")) or 
+    os.path.dirname(os.environ.get("AWS_CONFIG_FILE", "")) or
+    os.path.expanduser(os.path.join("~", ".aws"))
 )
 
-CREDENTIALS_CACHE_DIR = os.path.expanduser(
-    os.path.join("~", ".aws", "cli", "cache")
+SSO_TOKEN_DIR = os.environ.get("AWS_SSO_UTIL_SSO_TOKEN_DIR") or os.path.join(
+    AWS_CREDENTIALS_DIR,
+    "sso",
+    "cache",
+)
+
+CREDENTIALS_CACHE_DIR = os.environ.get("AWS_SSO_UTIL_CREDENTIALS_CACHE_DIR") or os.path.join(
+    AWS_CREDENTIALS_DIR,
+    "cli",
+    "cache",
 )
 
 LOGGER = logging.getLogger(__name__)