boto3-refresh-session 1.0.4__py3-none-any.whl → 6.2.5__py3-none-any.whl

boto3_refresh_session/methods/sts.py

@@ -0,0 +1,230 @@
+"""STS assume-role refreshable session implementation."""
+
+__all__ = ["STSRefreshableSession"]
+
+from typing import Callable
+
+from ..exceptions import BRSError, BRSWarning
+from ..utils import (
+    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 BRSError(
+                "'RoleArn' must be provided in 'assume_role_kwargs'!"
+            )
+
+        # 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"
+
+        # 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 BRSError(
+                "'mfa_token_provider' must be a callable that returns a "
+                "string representing an MFA token code!"
+            ) from err
+
+        # storing mfa_token_provider_kwargs
+        self.mfa_token_provider_kwargs = (
+            mfa_token_provider_kwargs if mfa_token_provider_kwargs else {}
+        )
+
+        # ensure SerialNumber is set appropriately with mfa_token_provider
+        if (
+            self.mfa_token_provider
+            and "SerialNumber" not in assume_role_kwargs
+        ):
+            raise BRSError(
+                "'SerialNumber' must be provided in 'assume_role_kwargs' "
+                "when using 'mfa_token_provider'!"
+            )
+
+        # 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 BRSError(
+                "'SerialNumber' and 'TokenCode' must be provided in "
+                "'assume_role_kwargs' when 'mfa_token_provider' is not set!"
+            )
+
+        # 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 BRSError(
+                    "'TokenCode' must be a 6-digit string per AWS MFA "
+                    "token specifications!"
+                )
+
+        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,148 +1,92 @@
+"""Public factory for constructing refreshable boto3 sessions."""
+
 from __future__ import annotations
 
-__doc__ = """
-A :class:`boto3.session.Session` object that automatically refreshes temporary
-credentials.
-"""
 __all__ = ["RefreshableSession"]
 
-from typing import Any, Dict
+from typing import get_args
+
+from .exceptions import BRSError
+from .utils import BaseRefreshableSession, Method
+
+
+class RefreshableSession:
+    """Factory class for constructing refreshable boto3 sessions using various
+    authentication methods, e.g. STS.
 
-from boto3 import client
-from boto3.session import Session
-from botocore.credentials import (
-    DeferredRefreshableCredentials,
-    RefreshableCredentials,
-)
+    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.
 
-class RefreshableSession(Session):
-    """Returns a :class:`boto3.session.Session` object with temporary credentials
-    that refresh automatically.
+    For additional information on required parameters, refer to the See Also
+    section below.
 
     Parameters
     ----------
-    assume_role_kwargs : dict
-        Required keyword arguments for the :meth:`STS.Client.assume_role` method.
+    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``.
-    sts_client_kwargs : dict, optional
-        Optional keyword arguments for the :class:`STS.Client` object. Default is
-        None.
+        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.
+    **kwargs : dict
+        Additional keyword arguments forwarded to the constructor of the
+        selected session class.
 
-    Notes
-    -----
-    Check the :ref:`authorization documentation <authorization>` for additional
-    information concerning how to authorize access to AWS.
-
-    Check the `AWS documentation <https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_temp.html>`_
-    for additional information concerning temporary security credentials in IAM.
-
-    Examples
+    See Also
     --------
-    In order to use this object, you are required to configure parameters for the
-    :meth:`STS.Client.assume_role` method.
-
-    >>> assume_role_kwargs = {
-    >>>     'RoleArn': '<your-role-arn>',
-    >>>     'RoleSessionName': '<your-role-session-name>',
-    >>>     'DurationSeconds': '<your-selection>',
-    >>>     ...
-    >>> }
-
-    You may also want to provide optional parameters for the :class:`STS.Client` object.
-
-    >>> sts_client_kwargs = {
-    >>>     ...
-    >>> }
-
-    You may also provide optional parameters for the :class:`boto3.session.Session` object
-    when initializing the ``RefreshableSession`` object. Below, we use the ``region_name``
-    parameter for illustrative purposes.
-
-    >>> session = boto3_refresh_session.RefreshableSession(
-    >>>     assume_role_kwargs=assume_role_kwargs,
-    >>>     sts_client_kwargs=sts_client_kwargs,
-    >>>     region_name='us-east-1',
-    >>> )
-
-    Using the ``session`` variable that you just created, you can now use all of the methods
-    available from the :class:`boto3.session.Session` object. In the below example, we
-    initialize an S3 client and list all available buckets.
-
-    >>> s3 = session.client(service_name='s3')
-    >>> buckets = s3.list_buckets()
-
-    There are two ways of refreshing temporary credentials automatically with the
-    ``RefreshableSession`` object: refresh credentials the moment they expire, or wait until
-    temporary credentials are explicitly needed. The latter is the default. The former must
-    be configured using the ``defer_refresh`` parameter, as shown below.
-
-    >>> session = boto3_refresh_session.RefreshableSession(
-    >>>     defer_refresh=False,
-    >>>     assume_role_kwargs=assume_role_kwargs,
-    >>>     sts_client_kwargs=sts_client_kwargs,
-    >>>     region_name='us-east-1',
-    >>> )
+    boto3_refresh_session.methods.custom.CustomRefreshableSession
+    boto3_refresh_session.methods.iot.x509.IOTX509RefreshableSession
+    boto3_refresh_session.methods.sts.STSRefreshableSession
     """
 
-    def __init__(
-        self,
-        assume_role_kwargs: Dict[Any],
-        defer_refresh: bool = True,
-        sts_client_kwargs: Dict[Any] = None,
-        **kwargs,
-    ):
-        # inheriting from boto3.session.Session
-        super().__init__(**kwargs)
-
-        # initializing custom parameters that are necessary outside of __init__
-        self.assume_role_kwargs = assume_role_kwargs
-
-        # initializing the STS client
-        if sts_client_kwargs is not None:
-            self._sts_client = client(service_name="sts", **sts_client_kwargs)
-        else:
-            self._sts_client = client(service_name="sts")
-
-        # determining how exactly to refresh expired temporary credentials
-        if not defer_refresh:
-            self._session._credentials = (
-                RefreshableCredentials.create_from_metadata(
-                    metadata=self._get_credentials(),
-                    refresh_using=self._get_credentials,
-                    method="sts-assume-role",
-                )
-            )
-        else:
-            self._session._credentials = DeferredRefreshableCredentials(
-                refresh_using=self._get_credentials, method="sts-assume-role"
+    def __new__(
+        cls, method: Method = "sts", **kwargs
+    ) -> BaseRefreshableSession:
+        if method not in (methods := cls.get_available_methods()):
+            raise BRSError(
+                f"{method!r} is an invalid method parameter. "
+                "Available methods are "
+                f"{', '.join(repr(meth) for meth in methods)}."
             )
 
-    def _get_credentials(self) -> Dict[Any]:
-        """Returns temporary credentials via AWS STS.
+        return BaseRefreshableSession.registry[method](**kwargs)
+
+    @classmethod
+    def get_available_methods(cls) -> list[str]:
+        """Lists all currently available credential refresh methods.
 
         Returns
         -------
-        dict
-            AWS temporary credentials.
+        list[str]
+            A list of all currently available credential refresh methods,
+            e.g. 'sts', 'custom'.
         """
 
-        # fetching temporary credentials
-        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"),
-            "token": temporary_credentials.get("SessionToken"),
-            "expiry_time": temporary_credentials.get("Expiration").isoformat(),
-        }
+        args = list(get_args(Method))
+        args.remove("__sentinel__")
+        return args

boto3_refresh_session/utils/__init__.py

@@ -0,0 +1,10 @@
+__all__ = []
+
+from . import cache, internal, typing
+from .cache import *
+from .internal import *
+from .typing import *
+
+__all__.extend(cache.__all__)
+__all__.extend(internal.__all__)
+__all__.extend(typing.__all__)

boto3_refresh_session/utils/cache.py

@@ -0,0 +1,94 @@
+"""Cache primitives for memoizing boto3 client instances.
+
+`ClientCache` provides a thread-safe mapping for cached clients and raises
+`BRSError` 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 BRSError
+
+
+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 BRSError(
+                    "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 BRSError("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 BRSError("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 BRSError("Client not found in cache.")
+            del self._cache[hash]
+            return client