# Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). You # may not use this file except in compliance with the License. A copy of # the License is located at # # http://aws.amazon.com/apache2.0/ # # or in the "license" file accompanying this file. This file is # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. import json import logging import os import threading from datetime import datetime, timedelta from typing import NamedTuple, Optional import dateutil.parser from dateutil.tz import tzutc from botocore import UNSIGNED from botocore.compat import total_seconds from botocore.config import Config from botocore.exceptions import ( ClientError, InvalidConfigError, TokenRetrievalError, ) from botocore.utils import CachedProperty, JSONFileCache, SSOTokenLoader logger = logging.getLogger(__name__) def _utc_now(): return datetime.now(tzutc()) def create_token_resolver(session): providers = [ SSOTokenProvider(session), ] return TokenProviderChain(providers=providers) def _serialize_utc_timestamp(obj): if isinstance(obj, datetime): return obj.strftime("%Y-%m-%dT%H:%M:%SZ") return obj def _sso_json_dumps(obj): return json.dumps(obj, default=_serialize_utc_timestamp) class FrozenAuthToken(NamedTuple): token: str expiration: Optional[datetime] = None class DeferredRefreshableToken: # The time at which we'll attempt to refresh, but not block if someone else # is refreshing. _advisory_refresh_timeout = 15 * 60 # The time at which all threads will block waiting for a refreshed token _mandatory_refresh_timeout = 10 * 60 # Refresh at most once every minute to avoid blocking every request _attempt_timeout = 60 def __init__(self, method, refresh_using, time_fetcher=_utc_now): self._time_fetcher = time_fetcher self._refresh_using = refresh_using self.method = method # The frozen token is protected by this lock self._refresh_lock = threading.Lock() self._frozen_token = None self._next_refresh = None def get_frozen_token(self): self._refresh() return self._frozen_token def _refresh(self): # If we don't need to refresh just return refresh_type = self._should_refresh() if not refresh_type: return None # Block for refresh if we're in the mandatory refresh window block_for_refresh = refresh_type == "mandatory" if self._refresh_lock.acquire(block_for_refresh): try: self._protected_refresh() finally: self._refresh_lock.release() def _protected_refresh(self): # This should only be called after acquiring the refresh lock # Another thread may have already refreshed, double check refresh refresh_type = self._should_refresh() if not refresh_type: return None try: now = self._time_fetcher() self._next_refresh = now + timedelta(seconds=self._attempt_timeout) self._frozen_token = self._refresh_using() except Exception: logger.warning( "Refreshing token failed during the %s refresh period.", refresh_type, exc_info=True, ) if refresh_type == "mandatory": # This refresh was mandatory, error must be propagated back raise if self._is_expired(): # Fresh credentials should never be expired raise TokenRetrievalError( provider=self.method, error_msg="Token has expired and refresh failed", ) def _is_expired(self): if self._frozen_token is None: return False expiration = self._frozen_token.expiration remaining = total_seconds(expiration - self._time_fetcher()) return remaining <= 0 def _should_refresh(self): if self._frozen_token is None: # We don't have a token yet, mandatory refresh return "mandatory" expiration = self._frozen_token.expiration if expiration is None: # No expiration, so assume we don't need to refresh. return None now = self._time_fetcher() if now < self._next_refresh: return None remaining = total_seconds(expiration - now) if remaining < self._mandatory_refresh_timeout: return "mandatory" elif remaining < self._advisory_refresh_timeout: return "advisory" return None class TokenProviderChain: def __init__(self, providers=None): if providers is None: providers = [] self._providers = providers def load_token(self): for provider in self._providers: token = provider.load_token() if token is not None: return token return None class SSOTokenProvider: METHOD = "sso" _REFRESH_WINDOW = 15 * 60 _SSO_TOKEN_CACHE_DIR = os.path.expanduser( os.path.join("~", ".aws", "sso", "cache") ) _SSO_CONFIG_VARS = [ "sso_start_url", "sso_region", ] _GRANT_TYPE = "refresh_token" DEFAULT_CACHE_CLS = JSONFileCache def __init__(self, session, cache=None, time_fetcher=_utc_now): self._session = session if cache is None: cache = self.DEFAULT_CACHE_CLS( self._SSO_TOKEN_CACHE_DIR, dumps_func=_sso_json_dumps, ) self._now = time_fetcher self._cache = cache self._token_loader = SSOTokenLoader(cache=self._cache) def _load_sso_config(self): loaded_config = self._session.full_config profiles = loaded_config.get("profiles", {}) sso_sessions = loaded_config.get("sso_sessions", {}) profile_name = self._session.get_config_variable("profile") if not profile_name: profile_name = "default" profile_config = profiles.get(profile_name, {}) if "sso_session" not in profile_config: return sso_session_name = profile_config["sso_session"] sso_config = sso_sessions.get(sso_session_name, None) if not sso_config: error_msg = ( f'The profile "{profile_name}" is configured to use the SSO ' f'token provider but the "{sso_session_name}" sso_session ' f"configuration does not exist." ) raise InvalidConfigError(error_msg=error_msg) missing_configs = [] for var in self._SSO_CONFIG_VARS: if var not in sso_config: missing_configs.append(var) if missing_configs: error_msg = ( f'The profile "{profile_name}" is configured to use the SSO ' f"token provider but is missing the following configuration: " f"{missing_configs}." ) raise InvalidConfigError(error_msg=error_msg) return { "session_name": sso_session_name, "sso_region": sso_config["sso_region"], "sso_start_url": sso_config["sso_start_url"], } @CachedProperty def _sso_config(self): return self._load_sso_config() @CachedProperty def _client(self): config = Config( region_name=self._sso_config["sso_region"], signature_version=UNSIGNED, ) return self._session.create_client("sso-oidc", config=config) def _attempt_create_token(self, token): response = self._client.create_token( grantType=self._GRANT_TYPE, clientId=token["clientId"], clientSecret=token["clientSecret"], refreshToken=token["refreshToken"], ) expires_in = timedelta(seconds=response["expiresIn"]) new_token = { "startUrl": self._sso_config["sso_start_url"], "region": self._sso_config["sso_region"], "accessToken": response["accessToken"], "expiresAt": self._now() + expires_in, # Cache the registration alongside the token "clientId": token["clientId"], "clientSecret": token["clientSecret"], "registrationExpiresAt": token["registrationExpiresAt"], } if "refreshToken" in response: new_token["refreshToken"] = response["refreshToken"] logger.info("SSO Token refresh succeeded") return new_token def _refresh_access_token(self, token): keys = ( "refreshToken", "clientId", "clientSecret", "registrationExpiresAt", ) missing_keys = [k for k in keys if k not in token] if missing_keys: msg = f"Unable to refresh SSO token: missing keys: {missing_keys}" logger.info(msg) return None expiry = dateutil.parser.parse(token["registrationExpiresAt"]) if total_seconds(expiry - self._now()) <= 0: logger.info(f"SSO token registration expired at {expiry}") return None try: return self._attempt_create_token(token) except ClientError: logger.warning("SSO token refresh attempt failed", exc_info=True) return None def _refresher(self): start_url = self._sso_config["sso_start_url"] session_name = self._sso_config["session_name"] logger.info(f"Loading cached SSO token for {session_name}") token_dict = self._token_loader(start_url, session_name=session_name) expiration = dateutil.parser.parse(token_dict["expiresAt"]) logger.debug(f"Cached SSO token expires at {expiration}") remaining = total_seconds(expiration - self._now()) if remaining < self._REFRESH_WINDOW: new_token_dict = self._refresh_access_token(token_dict) if new_token_dict is not None: token_dict = new_token_dict expiration = token_dict["expiresAt"] self._token_loader.save_token( start_url, token_dict, session_name=session_name ) return FrozenAuthToken( token_dict["accessToken"], expiration=expiration ) def load_token(self): if self._sso_config is None: return None return DeferredRefreshableToken( self.METHOD, self._refresher, time_fetcher=self._now )