codekingpro/portable-devtools

1# Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License"). You4# may not use this file except in compliance with the License. A copy of5# the License is located at6#7# http://aws.amazon.com/apache2.0/8#9# or in the "license" file accompanying this file. This file is10# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF11# ANY KIND, either express or implied. See the License for the specific12# language governing permissions and limitations under the License.13import json14import logging15import os16import threading17from datetime import datetime, timedelta18from typing import NamedTuple, Optional19 20import dateutil.parser21from dateutil.tz import tzutc22 23from botocore import UNSIGNED24from botocore.compat import total_seconds25from botocore.config import Config26from botocore.exceptions import (27    ClientError,28    InvalidConfigError,29    TokenRetrievalError,30)31from botocore.utils import CachedProperty, JSONFileCache, SSOTokenLoader32 33logger = logging.getLogger(__name__)34 35 36def _utc_now():37    return datetime.now(tzutc())38 39 40def create_token_resolver(session):41    providers = [42        SSOTokenProvider(session),43    ]44    return TokenProviderChain(providers=providers)45 46 47def _serialize_utc_timestamp(obj):48    if isinstance(obj, datetime):49        return obj.strftime("%Y-%m-%dT%H:%M:%SZ")50    return obj51 52 53def _sso_json_dumps(obj):54    return json.dumps(obj, default=_serialize_utc_timestamp)55 56 57class FrozenAuthToken(NamedTuple):58    token: str59    expiration: Optional[datetime] = None60 61 62class DeferredRefreshableToken:63    # The time at which we'll attempt to refresh, but not block if someone else64    # is refreshing.65    _advisory_refresh_timeout = 15 * 6066    # The time at which all threads will block waiting for a refreshed token67    _mandatory_refresh_timeout = 10 * 6068    # Refresh at most once every minute to avoid blocking every request69    _attempt_timeout = 6070 71    def __init__(self, method, refresh_using, time_fetcher=_utc_now):72        self._time_fetcher = time_fetcher73        self._refresh_using = refresh_using74        self.method = method75 76        # The frozen token is protected by this lock77        self._refresh_lock = threading.Lock()78        self._frozen_token = None79        self._next_refresh = None80 81    def get_frozen_token(self):82        self._refresh()83        return self._frozen_token84 85    def _refresh(self):86        # If we don't need to refresh just return87        refresh_type = self._should_refresh()88        if not refresh_type:89            return None90 91        # Block for refresh if we're in the mandatory refresh window92        block_for_refresh = refresh_type == "mandatory"93        if self._refresh_lock.acquire(block_for_refresh):94            try:95                self._protected_refresh()96            finally:97                self._refresh_lock.release()98 99    def _protected_refresh(self):100        # This should only be called after acquiring the refresh lock101        # Another thread may have already refreshed, double check refresh102        refresh_type = self._should_refresh()103        if not refresh_type:104            return None105 106        try:107            now = self._time_fetcher()108            self._next_refresh = now + timedelta(seconds=self._attempt_timeout)109            self._frozen_token = self._refresh_using()110        except Exception:111            logger.warning(112                "Refreshing token failed during the %s refresh period.",113                refresh_type,114                exc_info=True,115            )116            if refresh_type == "mandatory":117                # This refresh was mandatory, error must be propagated back118                raise119 120        if self._is_expired():121            # Fresh credentials should never be expired122            raise TokenRetrievalError(123                provider=self.method,124                error_msg="Token has expired and refresh failed",125            )126 127    def _is_expired(self):128        if self._frozen_token is None:129            return False130 131        expiration = self._frozen_token.expiration132        remaining = total_seconds(expiration - self._time_fetcher())133        return remaining <= 0134 135    def _should_refresh(self):136        if self._frozen_token is None:137            # We don't have a token yet, mandatory refresh138            return "mandatory"139 140        expiration = self._frozen_token.expiration141        if expiration is None:142            # No expiration, so assume we don't need to refresh.143            return None144 145        now = self._time_fetcher()146        if now < self._next_refresh:147            return None148 149        remaining = total_seconds(expiration - now)150 151        if remaining < self._mandatory_refresh_timeout:152            return "mandatory"153        elif remaining < self._advisory_refresh_timeout:154            return "advisory"155 156        return None157 158 159class TokenProviderChain:160    def __init__(self, providers=None):161        if providers is None:162            providers = []163        self._providers = providers164 165    def load_token(self):166        for provider in self._providers:167            token = provider.load_token()168            if token is not None:169                return token170        return None171 172 173class SSOTokenProvider:174    METHOD = "sso"175    _REFRESH_WINDOW = 15 * 60176    _SSO_TOKEN_CACHE_DIR = os.path.expanduser(177        os.path.join("~", ".aws", "sso", "cache")178    )179    _SSO_CONFIG_VARS = [180        "sso_start_url",181        "sso_region",182    ]183    _GRANT_TYPE = "refresh_token"184    DEFAULT_CACHE_CLS = JSONFileCache185 186    def __init__(187        self, session, cache=None, time_fetcher=_utc_now, profile_name=None188    ):189        self._session = session190        if cache is None:191            cache = self.DEFAULT_CACHE_CLS(192                self._SSO_TOKEN_CACHE_DIR,193                dumps_func=_sso_json_dumps,194            )195        self._now = time_fetcher196        self._cache = cache197        self._token_loader = SSOTokenLoader(cache=self._cache)198        self._profile_name = (199            profile_name200            or self._session.get_config_variable("profile")201            or 'default'202        )203 204    def _load_sso_config(self):205        loaded_config = self._session.full_config206        profiles = loaded_config.get("profiles", {})207        sso_sessions = loaded_config.get("sso_sessions", {})208        profile_config = profiles.get(self._profile_name, {})209 210        if "sso_session" not in profile_config:211            return212 213        sso_session_name = profile_config["sso_session"]214        sso_config = sso_sessions.get(sso_session_name, None)215 216        if not sso_config:217            error_msg = (218                f'The profile "{self._profile_name}" is configured to use the SSO '219                f'token provider but the "{sso_session_name}" sso_session '220                f"configuration does not exist."221            )222            raise InvalidConfigError(error_msg=error_msg)223 224        missing_configs = []225        for var in self._SSO_CONFIG_VARS:226            if var not in sso_config:227                missing_configs.append(var)228 229        if missing_configs:230            error_msg = (231                f'The profile "{self._profile_name}" is configured to use the SSO '232                f"token provider but is missing the following configuration: "233                f"{missing_configs}."234            )235            raise InvalidConfigError(error_msg=error_msg)236 237        return {238            "session_name": sso_session_name,239            "sso_region": sso_config["sso_region"],240            "sso_start_url": sso_config["sso_start_url"],241        }242 243    @CachedProperty244    def _sso_config(self):245        return self._load_sso_config()246 247    @CachedProperty248    def _client(self):249        config = Config(250            region_name=self._sso_config["sso_region"],251            signature_version=UNSIGNED,252        )253        return self._session.create_client("sso-oidc", config=config)254 255    def _attempt_create_token(self, token):256        response = self._client.create_token(257            grantType=self._GRANT_TYPE,258            clientId=token["clientId"],259            clientSecret=token["clientSecret"],260            refreshToken=token["refreshToken"],261        )262        expires_in = timedelta(seconds=response["expiresIn"])263        new_token = {264            "startUrl": self._sso_config["sso_start_url"],265            "region": self._sso_config["sso_region"],266            "accessToken": response["accessToken"],267            "expiresAt": self._now() + expires_in,268            # Cache the registration alongside the token269            "clientId": token["clientId"],270            "clientSecret": token["clientSecret"],271            "registrationExpiresAt": token["registrationExpiresAt"],272        }273        if "refreshToken" in response:274            new_token["refreshToken"] = response["refreshToken"]275        logger.info("SSO Token refresh succeeded")276        return new_token277 278    def _refresh_access_token(self, token):279        keys = (280            "refreshToken",281            "clientId",282            "clientSecret",283            "registrationExpiresAt",284        )285        missing_keys = [k for k in keys if k not in token]286        if missing_keys:287            msg = f"Unable to refresh SSO token: missing keys: {missing_keys}"288            logger.info(msg)289            return None290 291        expiry = dateutil.parser.parse(token["registrationExpiresAt"])292        if total_seconds(expiry - self._now()) <= 0:293            logger.info(f"SSO token registration expired at {expiry}")294            return None295 296        try:297            return self._attempt_create_token(token)298        except ClientError:299            logger.warning("SSO token refresh attempt failed", exc_info=True)300            return None301 302    def _refresher(self):303        start_url = self._sso_config["sso_start_url"]304        session_name = self._sso_config["session_name"]305        logger.info(f"Loading cached SSO token for {session_name}")306        token_dict = self._token_loader(start_url, session_name=session_name)307        expiration = dateutil.parser.parse(token_dict["expiresAt"])308        logger.debug(f"Cached SSO token expires at {expiration}")309 310        remaining = total_seconds(expiration - self._now())311        if remaining < self._REFRESH_WINDOW:312            new_token_dict = self._refresh_access_token(token_dict)313            if new_token_dict is not None:314                token_dict = new_token_dict315                expiration = token_dict["expiresAt"]316                self._token_loader.save_token(317                    start_url, token_dict, session_name=session_name318                )319 320        return FrozenAuthToken(321            token_dict["accessToken"], expiration=expiration322        )323 324    def load_token(self):325        if self._sso_config is None:326            return None327 328        return DeferredRefreshableToken(329            self.METHOD, self._refresher, time_fetcher=self._now330        )331