327 lines
11 KiB
Python
327 lines
11 KiB
Python
# 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
|
|
)
|