boto3-refresh-session 2.0.5__py3-none-any.whl → 7.1.3__py3-none-any.whl

boto3_refresh_session/methods/sts.py

@@ -1,75 +1,213 @@
-from __future__ import annotations
+# 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 Any
+from typing import Callable
 
-from ..exceptions import BRSWarning
-from ..session import BaseRefreshableSession
-from ..utils import AssumeRoleParams, STSClientParams, TemporaryCredentials
+from ..exceptions import BRSConfigurationError, BRSValidationError, BRSWarning
+from ..utils import (
+    AssumeRoleConfig,
+    AssumeRoleParams,
+    BaseRefreshableSession,
+    Identity,
+    STSClientConfig,
+    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
+    assume_role_kwargs : AssumeRoleParams | AssumeRoleConfig
         Required keyword arguments for :meth:`STS.Client.assume_role` (i.e.
-        boto3 STS client).
+        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 | STSClientConfig, 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``.
-    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.
+    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.
+
+    See Also
+    --------
+    boto3_refresh_session.utils.config.config.AssumeRoleConfig
+    boto3_refresh_session.utils.config.config.STSClientConfig
     """
 
     def __init__(
         self,
-        assume_role_kwargs: AssumeRoleParams,
-        defer_refresh: bool | None = None,
-        sts_client_kwargs: STSClientParams | None = None,
+        assume_role_kwargs: AssumeRoleParams | AssumeRoleConfig,
+        sts_client_kwargs: STSClientParams | STSClientConfig | None = None,
+        mfa_token_provider: Callable[[], str] | None = None,
+        mfa_token_provider_kwargs: dict | None = None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
-        self.assume_role_kwargs = 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(
-                    "'sts_client_kwargs' cannot contain values for "
-                    "'service_name'. Reverting to service_name = 'sts'."
+        # initializing asssume_role_kwargs attribute
+        match assume_role_kwargs:
+            case AssumeRoleConfig():
+                self.assume_role_kwargs = assume_role_kwargs
+            case _:
+                self.assume_role_kwargs = AssumeRoleConfig(
+                    **assume_role_kwargs
                 )
-                del sts_client_kwargs["service_name"]
-            self._sts_client = self.client(
-                service_name="sts", **sts_client_kwargs
+
+        # initializing sts_client_kwargs attribute
+        match sts_client_kwargs:
+            case STSClientConfig():
+                self.sts_client_kwargs = sts_client_kwargs
+            case None:
+                self.sts_client_kwargs = STSClientConfig()
+            case _:
+                self.sts_client_kwargs = STSClientConfig(**sts_client_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'."
             )
-        else:
-            self._sts_client = self.client(service_name="sts")
-
-        # mounting refreshable credentials
-        self.initialize(
-            credentials_method=self._get_credentials,
-            defer_refresh=defer_refresh is not False,
-            refresh_method="sts-assume-role",
+            del kwargs["refresh_method"]
+
+        # setting 'RoleSessionName' if not provided
+        self.assume_role_kwargs.RoleSessionName = self.assume_role_kwargs.get(
+            "RoleSessionName", "boto3-refresh-session"
         )
 
+        # 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 or {}
+
+        # ensure SerialNumber is set appropriately with mfa_token_provider
+        if (
+            self.mfa_token_provider
+            and self.assume_role_kwargs.SerialNumber is None
+        ):
+            raise BRSConfigurationError(
+                "'SerialNumber' must be provided in 'assume_role_kwargs' "
+                "when using 'mfa_token_provider'!",
+                param="SerialNumber",
+            )
+
+        # ensure SerialNumber and TokenCode are set in the absence of
+        # mfa_token_provider
+        if (
+            self.mfa_token_provider is None
+            and (
+                self.assume_role_kwargs.SerialNumber is not None
+                and self.assume_role_kwargs.TokenCode is None
+            )
+            or (
+                self.assume_role_kwargs.SerialNumber is None
+                and self.assume_role_kwargs.TokenCode is not None
+            )
+        ):
+            raise BRSConfigurationError(
+                "'SerialNumber' and 'TokenCode' must be provided in "
+                "'assume_role_kwargs' when 'mfa_token_provider' is not set "
+                "and 'SerialNumber' or 'TokenCode' is missing!",
+                param="SerialNumber/TokenCode",
+            )
+
+        # warn if TokenCode provided with mfa_token_provider
+        if (
+            self.mfa_token_provider
+            and self.assume_role_kwargs.TokenCode is not None
+        ):
+            BRSWarning.warn(
+                "'TokenCode' provided in 'assume_role_kwargs' will be "
+                "ignored and overridden by 'mfa_token_provider' on each "
+                "refresh."
+            )
+
+        # initializing BRSSession
+        super().__init__(refresh_method="sts-assume-role", **kwargs)
+
+        # initializing STS client attribute
+        self._sts_client = self.client(**self.sts_client_kwargs)
+
     def _get_credentials(self) -> TemporaryCredentials:
+        # override TokenCode with fresh token from provider if configured
+        if self.mfa_token_provider is not None:
+            self.assume_role_kwargs.TokenCode = self.mfa_token_provider(
+                **self.mfa_token_provider_kwargs
+            )
+
         temporary_credentials = self._sts_client.assume_role(
             **self.assume_role_kwargs
         )["Credentials"]
+
         return {
             "access_key": temporary_credentials.get("AccessKeyId"),
             "secret_key": temporary_credentials.get("SecretAccessKey"),
@@ -77,12 +215,12 @@ class STSRefreshableSession(BaseRefreshableSession, registry_key="sts"):
             "expiry_time": temporary_credentials.get("Expiration").isoformat(),
         }
 
-    def get_identity(self) -> dict[str, Any]:
+    def get_identity(self) -> Identity:
         """Returns metadata about the identity assumed.
 
         Returns
         -------
-        dict[str, Any]
+        Identity
             Dict containing caller identity according to AWS STS.
         """
 

boto3_refresh_session/session.py

@@ -1,37 +1,17 @@
+# 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."""
+
 from __future__ import annotations
 
 __all__ = ["RefreshableSession"]
 
 from typing import get_args
 
-from .exceptions import BRSError
-from .utils import BRSSession, CredentialProvider, Method, Registry
-
-
-class BaseRefreshableSession(
-    Registry[Method],
-    CredentialProvider,
-    BRSSession,
-    registry_key="__sentinel__",
-):
-    """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.
-    """
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
+from .exceptions import BRSValidationError
+from .utils import BaseRefreshableSession, Method
 
 
 class RefreshableSession:
@@ -52,6 +32,28 @@ class RefreshableSession:
     method : Method
         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
     ----------------
@@ -62,18 +64,32 @@ class RefreshableSession:
     See Also
     --------
     boto3_refresh_session.methods.custom.CustomRefreshableSession
+    boto3_refresh_session.methods.iot.x509.IOTX509RefreshableSession
     boto3_refresh_session.methods.sts.STSRefreshableSession
-    boto3_refresh_session.methods.ecs.ECSRefreshableSession
+
+    Examples
+    --------
+
+    Basic initialization using STS AssumeRole (i.e. ``method="sts"``):
+
+    >>> from boto3_refresh_session import AssumeRoleConfig, RefreshableSession
+    >>> session = RefreshableSession(
+    ...     AssumeRoleConfig(RoleArn="<your-role-arn>"),
+    ...     region_name="us-east-1"
+    ... )
+    >>> s3 = session.client("s3")
     """
 
     def __new__(
         cls, method: Method = "sts", **kwargs
     ) -> BaseRefreshableSession:
         if method not in (methods := cls.get_available_methods()):
-            raise BRSError(
+            raise BRSValidationError(
                 f"{method!r} is an invalid method parameter. "
                 "Available methods are "
-                f"{', '.join(repr(meth) for meth in methods)}."
+                f"{', '.join(repr(meth) for meth in methods)}.",
+                param="method",
+                value=method,
             )
 
         return BaseRefreshableSession.registry[method](**kwargs)
@@ -86,7 +102,7 @@ class RefreshableSession:
         -------
         list[str]
             A list of all currently available credential refresh methods,
-            e.g. 'sts', 'ecs', 'custom'.
+            e.g. 'sts', 'custom'.
         """
 
         args = list(get_args(Method))

boto3_refresh_session/utils/__init__.py

@@ -0,0 +1,18 @@
+# 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, config, constants, internal, typing
+from .cache import *
+from .config import *
+from .constants import *
+from .internal import *
+from .typing import *
+
+__all__.extend(cache.__all__)
+__all__.extend(config.__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/config/__init__.py

@@ -0,0 +1,10 @@
+# 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 config
+from .config import *
+
+__all__.extend(config.__all__)