boto3-refresh-session 1.3.11__py3-none-any.whl → 7.0.1__py3-none-any.whl

boto3_refresh_session/methods/sts.py

@@ -0,0 +1,276 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+"""STS assume-role refreshable session implementation."""
+
+__all__ = ["STSRefreshableSession"]
+
+from typing import Callable
+
+from ..exceptions import BRSConfigurationError, BRSValidationError, BRSWarning
+from ..utils import (
+    MFA_SERIAL_PATTERN,
+    ROLE_ARN_PATTERN,
+    ROLE_SESSION_NAME_PATTERN,
+    AssumeRoleParams,
+    BaseRefreshableSession,
+    Identity,
+    STSClientParams,
+    TemporaryCredentials,
+    refreshable_session,
+)
+
+
+@refreshable_session
+class STSRefreshableSession(BaseRefreshableSession, registry_key="sts"):
+    """A :class:`boto3.session.Session` object that automatically refreshes
+    temporary AWS credentials using an IAM role that is assumed via STS.
+
+    Parameters
+    ----------
+    assume_role_kwargs : AssumeRoleParams
+        Required keyword arguments for :meth:`STS.Client.assume_role` (i.e.
+        boto3 STS client). ``RoleArn`` is required. ``RoleSessionName`` will
+        default to 'boto3-refresh-session' if not provided.
+
+        For MFA authentication, two modalities are supported:
+
+        1. **Dynamic tokens (recommended)**: Provide ``SerialNumber`` in
+           ``assume_role_kwargs`` and pass ``mfa_token_provider`` callable.
+           The provider callable will be invoked on each refresh to obtain
+           fresh MFA tokens. Do not include ``TokenCode`` in this case.
+
+        2. **Static/injectable tokens**: Provide both ``SerialNumber`` and
+           ``TokenCode`` in ``assume_role_kwargs``. You are responsible for
+           updating ``assume_role_kwargs["TokenCode"]`` before the token
+           expires.
+    sts_client_kwargs : STSClientParams, optional
+        Optional keyword arguments for the :class:`STS.Client` object. Do not
+        provide values for ``service_name`` as they are unnecessary. Default
+        is None.
+    mfa_token_provider : Callable[[], str], optional
+        An optional callable that returns a string representing a fresh MFA
+        token code. If provided, this will be called during each credential
+        refresh to obtain a new token, which overrides any ``TokenCode`` in
+        ``assume_role_kwargs``. When using this parameter, ``SerialNumber``
+        must be provided in ``assume_role_kwargs``. Default is None.
+    mfa_token_provider_kwargs : dict, optional
+        Optional keyword arguments to pass to the ``mfa_token_provider``
+        callable. Default is None.
+    defer_refresh : bool, optional
+        If ``True`` then temporary credentials are not automatically refreshed
+        until they are explicitly needed. If ``False`` then temporary
+        credentials refresh immediately upon expiration. It is highly
+        recommended that you use ``True``. Default is ``True``.
+    advisory_timeout : int, optional
+        USE THIS ARGUMENT WITH CAUTION!!!
+
+        Botocore will attempt to refresh credentials early according to
+        this value (in seconds), but will continue using the existing
+        credentials if refresh fails. Default is 15 minutes (900 seconds).
+    mandatory_timeout : int, optional
+        USE THIS ARGUMENT WITH CAUTION!!!
+
+        Botocore requires a successful refresh before continuing. If
+        refresh fails in this window (in seconds), API calls may fail.
+        Default is 10 minutes (600 seconds).
+    cache_clients : bool, optional
+        If ``True`` then clients created by this session will be cached and
+        reused for subsequent calls to :meth:`client()` with the same
+        parameter signatures. Due to the memory overhead of clients, the
+        default is ``True`` in order to protect system resources.
+
+    Other Parameters
+    ----------------
+    kwargs : dict
+        Optional keyword arguments for the :class:`boto3.session.Session`
+        object.
+    """
+
+    def __init__(
+        self,
+        assume_role_kwargs: AssumeRoleParams,
+        sts_client_kwargs: STSClientParams | None = None,
+        mfa_token_provider: Callable[[], str] | None = None,
+        mfa_token_provider_kwargs: dict | None = None,
+        **kwargs,
+    ):
+        # ensuring 'refresh_method' is not set manually
+        if "refresh_method" in kwargs:
+            BRSWarning.warn(
+                "'refresh_method' cannot be set manually. "
+                "Reverting to 'sts-assume-role'."
+            )
+            del kwargs["refresh_method"]
+
+        # verifying 'RoleArn' is provided in 'assume_role_kwargs'
+        if "RoleArn" not in assume_role_kwargs:
+            raise BRSConfigurationError(
+                "'RoleArn' must be provided in 'assume_role_kwargs'!",
+                param="RoleArn",
+            )
+
+        # verifying 'RoleArn' format
+        if not ROLE_ARN_PATTERN.match(assume_role_kwargs["RoleArn"]):
+            raise BRSValidationError(
+                "'RoleArn' in 'assume_role_kwargs' is not a valid AWS "
+                "Role ARN!",
+                param="RoleArn",
+                value=assume_role_kwargs.get("RoleArn"),
+            )
+
+        # setting default 'RoleSessionName' if not provided
+        if "RoleSessionName" not in assume_role_kwargs:
+            BRSWarning.warn(
+                "'RoleSessionName' not provided in "
+                "'assume_role_kwargs'! Defaulting to "
+                "'boto3-refresh-session'."
+            )
+            assume_role_kwargs["RoleSessionName"] = "boto3-refresh-session"
+
+        # verifying 'RoleSessionName' format
+        if not ROLE_SESSION_NAME_PATTERN.match(
+            assume_role_kwargs["RoleSessionName"]
+        ):
+            raise BRSValidationError(
+                "'RoleSessionName' in 'assume_role_kwargs' is not valid! "
+                "It must be 2-64 characters long and can contain only "
+                "alphanumeric characters and the following symbols: "
+                "'+=,.@-'.",
+                param="RoleSessionName",
+                value=assume_role_kwargs.get("RoleSessionName"),
+            )
+
+        # store MFA token provider
+        try:
+            # verifying type of mfa_token_provider
+            assert (
+                isinstance(mfa_token_provider, Callable)
+                or mfa_token_provider is None
+            )
+            self.mfa_token_provider = mfa_token_provider
+        except AssertionError as err:
+            raise BRSValidationError(
+                "'mfa_token_provider' must be a callable that returns a "
+                "string representing an MFA token code!",
+                param="mfa_token_provider",
+            ) from err
+
+        # storing mfa_token_provider_kwargs
+        self.mfa_token_provider_kwargs = (
+            mfa_token_provider_kwargs if mfa_token_provider_kwargs else {}
+        )
+
+        # verifying 'SerialNumber' format if provided
+        if "SerialNumber" in assume_role_kwargs:
+            if not MFA_SERIAL_PATTERN.match(
+                assume_role_kwargs["SerialNumber"]
+            ):
+                raise BRSValidationError(
+                    "'SerialNumber' in 'assume_role_kwargs' is not a valid "
+                    "AWS MFA device ARN!",
+                    param="SerialNumber",
+                    value=assume_role_kwargs.get("SerialNumber"),
+                )
+
+        # ensure SerialNumber is set appropriately with mfa_token_provider
+        if (
+            self.mfa_token_provider
+            and "SerialNumber" not in assume_role_kwargs
+        ):
+            raise BRSConfigurationError(
+                "'SerialNumber' must be provided in 'assume_role_kwargs' "
+                "when using 'mfa_token_provider'!",
+                param="SerialNumber",
+            )
+
+        # ensure SerialNumber and TokenCode are set without mfa_token_provider
+        if (
+            self.mfa_token_provider is None
+            and (
+                "SerialNumber" in assume_role_kwargs
+                and "TokenCode" not in assume_role_kwargs
+            )
+            or (
+                "SerialNumber" not in assume_role_kwargs
+                and "TokenCode" in assume_role_kwargs
+            )
+        ):
+            raise BRSConfigurationError(
+                "'SerialNumber' and 'TokenCode' must be provided in "
+                "'assume_role_kwargs' when 'mfa_token_provider' is not set!",
+                param="SerialNumber/TokenCode",
+            )
+
+        # warn if TokenCode provided with mfa_token_provider
+        if self.mfa_token_provider and "TokenCode" in assume_role_kwargs:
+            BRSWarning.warn(
+                "'TokenCode' provided in 'assume_role_kwargs' will be "
+                "ignored and overridden by 'mfa_token_provider' on each "
+                "refresh."
+            )
+
+        # initializing assume role kwargs attribute
+        self.assume_role_kwargs = assume_role_kwargs
+
+        # initializing BRSSession
+        super().__init__(refresh_method="sts-assume-role", **kwargs)
+
+        if sts_client_kwargs is not None:
+            # overwriting 'service_name' if if appears in sts_client_kwargs
+            if "service_name" in sts_client_kwargs:
+                BRSWarning.warn(
+                    "'sts_client_kwargs' cannot contain values for "
+                    "'service_name'. Reverting to service_name = 'sts'."
+                )
+                del sts_client_kwargs["service_name"]
+            self._sts_client = self.client(
+                service_name="sts", **sts_client_kwargs
+            )
+        else:
+            self._sts_client = self.client(service_name="sts")
+
+    def _get_credentials(self) -> TemporaryCredentials:
+        params = dict(self.assume_role_kwargs)
+
+        # override TokenCode with fresh token from provider if configured
+        if self.mfa_token_provider:
+            params["TokenCode"] = self.mfa_token_provider(
+                **self.mfa_token_provider_kwargs
+            )
+
+        # validating TokenCode format
+        if (token_code := params.get("TokenCode")) is not None:
+            if (
+                not isinstance(token_code, str)
+                or len(token_code) != 6
+                or not token_code.isdigit()
+            ):
+                raise BRSValidationError(
+                    "'TokenCode' must be a 6-digit string per AWS MFA "
+                    "token specifications!",
+                    param="TokenCode",
+                )
+
+        temporary_credentials = self._sts_client.assume_role(**params)[
+            "Credentials"
+        ]
+
+        return {
+            "access_key": temporary_credentials.get("AccessKeyId"),
+            "secret_key": temporary_credentials.get("SecretAccessKey"),
+            "token": temporary_credentials.get("SessionToken"),
+            "expiry_time": temporary_credentials.get("Expiration").isoformat(),
+        }
+
+    def get_identity(self) -> Identity:
+        """Returns metadata about the identity assumed.
+
+        Returns
+        -------
+        Identity
+            Dict containing caller identity according to AWS STS.
+        """
+
+        return self._sts_client.get_caller_identity()

boto3_refresh_session/session.py

@@ -1,151 +1,86 @@
-from __future__ import annotations
-
-__all__ = ["RefreshableSession"]
-
-from abc import ABC, abstractmethod
-from typing import Any, Callable, ClassVar, Literal, get_args
-
-from boto3.session import Session
-from botocore.credentials import (
-    DeferredRefreshableCredentials,
-    RefreshableCredentials,
-)
-
-from .exceptions import BRSError, BRSWarning
-
-#: Type alias for all currently available credential refresh methods.
-Method = Literal["sts", "ecs", "custom"]
-RefreshMethod = Literal["sts-assume-role", "ecs-container-metadata", "custom"]
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at https://mozilla.org/MPL/2.0/.
 
+"""Public factory for constructing refreshable boto3 sessions."""
 
-class BaseRefreshableSession(ABC, Session):
-    """Abstract base class for implementing refreshable AWS sessions.
-
-    Provides a common interface and factory registration mechanism
-    for subclasses that generate temporary credentials using various
-    AWS authentication methods (e.g., STS).
-
-    Subclasses must implement ``_get_credentials()`` and ``get_identity()``.
-    They should also register themselves using the ``method=...`` argument
-    to ``__init_subclass__``.
-
-    Parameters
-    ----------
-    registry : dict[str, type[BaseRefreshableSession]]
-        Class-level registry mapping method names to registered session types.
-    """
-
-    # adding this and __init_subclass__ to avoid circular imports
-    # as well as simplify future addition of new methods
-    registry: ClassVar[dict[Method, type[BaseRefreshableSession]]] = {}
-
-    def __init_subclass__(cls, method: Method):
-        super().__init_subclass__()
-
-        # guarantees that methods are unique
-        if method in BaseRefreshableSession.registry:
-            BRSWarning(
-                f"Method {repr(method)} is already registered. Overwriting."
-            )
-
-        BaseRefreshableSession.registry[method] = cls
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    @abstractmethod
-    def _get_credentials(self) -> dict[str, str]: ...
-
-    @abstractmethod
-    def get_identity(self) -> dict[str, Any]: ...
-
-    def _refresh_using(
-        self,
-        credentials_method: Callable,
-        defer_refresh: bool,
-        refresh_method: RefreshMethod,
-    ):
-        # determining how exactly to refresh expired temporary credentials
-        if not defer_refresh:
-            self._credentials = RefreshableCredentials.create_from_metadata(
-                metadata=credentials_method(),
-                refresh_using=credentials_method,
-                method=refresh_method,
-            )
-        else:
-            self._credentials = DeferredRefreshableCredentials(
-                refresh_using=credentials_method, method=refresh_method
-            )
-
-    def refreshable_credentials(self) -> dict[str, str]:
-        """The current temporary AWS security credentials.
-
-        Returns
-        -------
-        dict[str, str]
-            Temporary AWS security credentials containing:
-                AWS_ACCESS_KEY_ID : str
-                    AWS access key identifier.
-                AWS_SECRET_ACCESS_KEY : str
-                    AWS secret access key.
-                AWS_SESSION_TOKEN : str
-                    AWS session token.
-        """
+from __future__ import annotations
 
-        creds = self.get_credentials().get_frozen_credentials()
-        return {
-            "AWS_ACCESS_KEY_ID": creds.access_key,
-            "AWS_SECRET_ACCESS_KEY": creds.secret_key,
-            "AWS_SESSION_TOKEN": creds.token,
-        }
+__all__ = ["RefreshableSession"]
 
-    @property
-    def credentials(self) -> dict[str, str]:
-        """The current temporary AWS security credentials."""
+from typing import get_args
 
-        return self.refreshable_credentials()
+from .exceptions import BRSValidationError
+from .utils import BaseRefreshableSession, Method
 
 
 class RefreshableSession:
-    """Factory class for constructing refreshable boto3 sessions using various authentication
-    methods, e.g. STS.
+    """Factory class for constructing refreshable boto3 sessions using various
+    authentication methods, e.g. STS.
 
-    This class provides a unified interface for creating boto3 sessions whose credentials are
-    automatically refreshed in the background.
+    This class provides a unified interface for creating boto3 sessions whose
+    credentials are automatically refreshed in the background.
 
-    Use ``RefreshableSession(method="...")`` to construct an instance using the desired method.
+    Use ``RefreshableSession(method="...")`` to construct an instance using
+    the desired method.
 
-    For additional information on required parameters, refer to the See Also section below.
+    For additional information on required parameters, refer to the See Also
+    section below.
 
     Parameters
     ----------
     method : Method
-        The authentication and refresh method to use for the session. Must match a registered method name.
-        Default is "sts".
+        The authentication and refresh method to use for the session. Must
+        match a registered method name. Default is "sts".
+    defer_refresh : bool, optional
+        If ``True`` then temporary credentials are not automatically refreshed
+        until they are explicitly needed. If ``False`` then temporary
+        credentials refresh immediately upon expiration. It is highly
+        recommended that you use ``True``. Default is ``True``.
+    advisory_timeout : int, optional
+        USE THIS ARGUMENT WITH CAUTION!!!
+
+        Botocore will attempt to refresh credentials early according to
+        this value (in seconds), but will continue using the existing
+        credentials if refresh fails. Default is 15 minutes (900 seconds).
+    mandatory_timeout : int, optional
+        USE THIS ARGUMENT WITH CAUTION!!!
+
+        Botocore requires a successful refresh before continuing. If
+        refresh fails in this window (in seconds), API calls may fail.
+        Default is 10 minutes (600 seconds).
+    cache_clients : bool, optional
+        If ``True`` then clients created by this session will be cached and
+        reused for subsequent calls to :meth:`client()` with the same
+        parameter signatures. Due to the memory overhead of clients, the
+        default is ``True`` in order to protect system resources.
 
     Other Parameters
     ----------------
     **kwargs : dict
-        Additional keyword arguments forwarded to the constructor of the selected session class.
+        Additional keyword arguments forwarded to the constructor of the
+        selected session class.
 
     See Also
     --------
-    boto3_refresh_session.custom.CustomRefreshableSession
-    boto3_refresh_session.sts.STSRefreshableSession
-    boto3_refresh_session.ecs.ECSRefreshableSession
+    boto3_refresh_session.methods.custom.CustomRefreshableSession
+    boto3_refresh_session.methods.iot.x509.IOTX509RefreshableSession
+    boto3_refresh_session.methods.sts.STSRefreshableSession
     """
 
     def __new__(
         cls, method: Method = "sts", **kwargs
     ) -> BaseRefreshableSession:
         if method not in (methods := cls.get_available_methods()):
-            raise BRSError(
-                f"{repr(method)} is an invalid method parameter. Available methods are "
-                f"{', '.join(repr(meth) for meth in methods)}."
+            raise BRSValidationError(
+                f"{method!r} is an invalid method parameter. "
+                "Available methods are "
+                f"{', '.join(repr(meth) for meth in methods)}.",
+                param="method",
+                value=method,
             )
 
-        obj = BaseRefreshableSession.registry[method]
-        return obj(**kwargs)
+        return BaseRefreshableSession.registry[method](**kwargs)
 
     @classmethod
     def get_available_methods(cls) -> list[str]:
@@ -154,7 +89,10 @@ class RefreshableSession:
         Returns
         -------
         list[str]
-            A list of all currently available credential refresh methods, e.g. 'sts'.
+            A list of all currently available credential refresh methods,
+            e.g. 'sts', 'custom'.
         """
 
-        return list(get_args(Method))
+        args = list(get_args(Method))
+        args.remove("__sentinel__")
+        return args

boto3_refresh_session/utils/__init__.py

@@ -0,0 +1,16 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+__all__ = []
+
+from . import cache, constants, internal, typing
+from .cache import *
+from .constants import *
+from .internal import *
+from .typing import *
+
+__all__.extend(cache.__all__)
+__all__.extend(constants.__all__)
+__all__.extend(internal.__all__)
+__all__.extend(typing.__all__)

boto3_refresh_session/utils/cache.py

@@ -0,0 +1,98 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+"""Cache primitives for memoizing boto3 client instances.
+
+`ClientCache` provides a thread-safe mapping for cached clients and raises
+`BRSCacheError` when lookups or mutations violate the expected cache contract.
+"""
+
+__all__ = ["ClientCache"]
+
+from threading import Lock
+from typing import Hashable, Optional
+
+from botocore.client import BaseClient
+
+from ..exceptions import BRSCacheExistsError, BRSCacheNotFoundError
+
+
+class ClientCache:
+    """A thread-safe cache for storing boto3 clients which can be used like a
+    dictionary."""
+
+    def __init__(self):
+        self._cache: dict[Hashable, BaseClient] = {}
+        self._lock = Lock()
+
+    def __len__(self) -> int:
+        with self._lock:
+            return len(self._cache)
+
+    def __contains__(self, hash: Hashable) -> bool:
+        with self._lock:
+            return hash in self._cache
+
+    def __iter__(self):
+        with self._lock:
+            return iter(self._cache.keys())
+
+    def __getitem__(self, hash: Hashable) -> BaseClient:
+        with self._lock:
+            try:
+                return self._cache[hash]
+            except KeyError as err:
+                raise BRSCacheNotFoundError(
+                    "The client you requested has not been cached."
+                ) from err
+
+    def __setitem__(self, hash: Hashable, client: BaseClient) -> None:
+        with self._lock:
+            if hash in self._cache:
+                raise BRSCacheExistsError("Client already exists in cache.")
+
+            self._cache[hash] = client
+
+    def __delitem__(self, hash: Hashable) -> None:
+        with self._lock:
+            if hash not in self._cache:
+                raise BRSCacheNotFoundError("Client not found in cache.")
+            del self._cache[hash]
+
+    def keys(self) -> tuple[Hashable, ...]:
+        """Returns the keys in the cache."""
+
+        with self._lock:
+            return tuple(self._cache.keys())
+
+    def values(self) -> tuple[BaseClient, ...]:
+        """Returns the clients from the cache."""
+
+        with self._lock:
+            return tuple(self._cache.values())
+
+    def items(self) -> tuple[tuple[Hashable, BaseClient], ...]:
+        """Returns the items in the cache as (hash, BaseClient) tuples."""
+
+        with self._lock:
+            return tuple(self._cache.items())
+
+    def get(
+        self, hash: Hashable, default: Optional[BaseClient] = None
+    ) -> Optional[BaseClient]:
+        """Gets the client using the given signature, or returns None if no
+        default is provided.
+        """
+
+        with self._lock:
+            return self._cache.get(hash, default)
+
+    def pop(self, hash: Hashable) -> BaseClient:
+        """Pops and returns the client using the given signature."""
+
+        with self._lock:
+            if (client := self._cache.get(hash)) is None:
+                raise BRSCacheNotFoundError("Client not found in cache.")
+            del self._cache[hash]
+            return client

boto3_refresh_session/utils/constants.py

@@ -0,0 +1,16 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+__all__ = [
+    "ROLE_ARN_PATTERN",
+    "MFA_SERIAL_PATTERN",
+    "ROLE_SESSION_NAME_PATTERN",
+]
+
+from re import compile
+
+# AWS ARN validation patterns
+ROLE_ARN_PATTERN = compile(r"^arn:aws[a-z-]*:iam::\d{12}:role/[\w+=,.@-]+$")
+MFA_SERIAL_PATTERN = compile(r"^arn:aws[a-z-]*:iam::\d{12}:mfa/[\w+=,.@-]+$")
+ROLE_SESSION_NAME_PATTERN = compile(r"^[a-zA-Z0-9+=,.@-]{2,64}$")