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

boto3_refresh_session/methods/sts.py

@@ -1,14 +1,28 @@
-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 (
+    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.
@@ -17,16 +31,55 @@ class STSRefreshableSession(BaseRefreshableSession, registry_key="sts"):
     ----------
     assume_role_kwargs : AssumeRoleParams
         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, 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
     ----------------
@@ -38,17 +91,136 @@ class STSRefreshableSession(BaseRefreshableSession, registry_key="sts"):
     def __init__(
         self,
         assume_role_kwargs: AssumeRoleParams,
-        defer_refresh: bool | None = None,
         sts_client_kwargs: STSClientParams | None = None,
+        mfa_token_provider: Callable[[], str] | None = None,
+        mfa_token_provider_kwargs: dict | None = None,
         **kwargs,
     ):
-        super().__init__(**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(
+                BRSWarning.warn(
                     "'sts_client_kwargs' cannot contain values for "
                     "'service_name'. Reverting to service_name = 'sts'."
                 )
@@ -59,17 +231,32 @@ class STSRefreshableSession(BaseRefreshableSession, registry_key="sts"):
         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",
-        )
-
     def _get_credentials(self) -> TemporaryCredentials:
-        temporary_credentials = self._sts_client.assume_role(
-            **self.assume_role_kwargs
-        )["Credentials"]
+        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"),
@@ -77,12 +264,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,20 @@ 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
     """
 
     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 +90,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,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}$")