[minor] Adding ECS module to boto3-refresh-session (#57)

* ecs

* ecs changes

* readme updates

* updated _get_credentials

* docs

* doc updates:

* static method

* custom exceptions + repr

* errors docs updates

* flake8 updates

* black updates

Author: michaelthomasletts

README.md

@@ -54,6 +54,7 @@
 
 - Auto-refreshing credentials for long-lived `boto3` sessions
 - Drop-in replacement for `boto3.session.Session`
+- Supports automatic refresh methods for STS and ECS
 - Supports `assume_role` configuration, custom STS clients, and profile / region configuration, as well as all other parameters supported by `boto3.session.Session`
 - Tested, documented, and published to PyPI
 

README.template.md

@@ -54,6 +54,7 @@
 
 - Auto-refreshing credentials for long-lived `boto3` sessions
 - Drop-in replacement for `boto3.session.Session`
+- Supports automatic refresh methods for STS and ECS
 - Supports `assume_role` configuration, custom STS clients, and profile / region configuration, as well as all other parameters supported by `boto3.session.Session`
 - Tested, documented, and published to PyPI
 

boto3_refresh_session/__init__.py

@@ -1,3 +1,4 @@
+from .ecs import ECSRefreshableSession
 from .session import RefreshableSession
 from .sts import STSRefreshableSession
 

boto3_refresh_session/ecs.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+__all__ = ["ECSRefreshableSession"]
+
+import os
+
+import requests
+
+from .exceptions import BRSError
+from .session import BaseRefreshableSession
+
+_ECS_CREDENTIALS_RELATIVE_URI = "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"
+_ECS_CREDENTIALS_FULL_URI = "AWS_CONTAINER_CREDENTIALS_FULL_URI"
+_ECS_AUTHORIZATION_TOKEN = "AWS_CONTAINER_AUTHORIZATION_TOKEN"
+_DEFAULT_ENDPOINT_BASE = "http://169.254.170.2"
+
+
+class ECSRefreshableSession(BaseRefreshableSession, method="ecs"):
+    """A boto3 session that automatically refreshes temporary AWS credentials
+    from the ECS container credentials metadata endpoint.
+
+    Parameters
+    ----------
+    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``.
+
+    Other Parameters
+    ----------------
+    kwargs : dict
+        Optional keyword arguments passed to :class:`boto3.session.Session`.
+    """
+
+    def __init__(self, defer_refresh: bool | None = None, **kwargs):
+        super().__init__(**kwargs)
+
+        self._endpoint = self._resolve_endpoint()
+        self._headers = self._build_headers()
+        self._http = self._init_http_session()
+
+        self._refresh_using(
+            credentials_method=self._get_credentials,
+            defer_refresh=defer_refresh is not False,
+            refresh_method="ecs-container-metadata",
+        )
+
+    def _resolve_endpoint(self) -> str:
+        uri = os.environ.get(_ECS_CREDENTIALS_FULL_URI) or os.environ.get(
+            _ECS_CREDENTIALS_RELATIVE_URI
+        )
+        if not uri:
+            raise BRSError(
+                "Neither AWS_CONTAINER_CREDENTIALS_FULL_URI nor "
+                "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI is set. "
+                "Are you running inside an ECS container?"
+            )
+        if uri.startswith("http://") or uri.startswith("https://"):
+            return uri
+        return f"{_DEFAULT_ENDPOINT_BASE}{uri}"
+
+    def _build_headers(self) -> dict[str, str]:
+        token = os.environ.get(_ECS_AUTHORIZATION_TOKEN)
+        if token:
+            return {"Authorization": f"Bearer {token}"}
+        return {}
+
+    def _init_http_session(self) -> requests.Session:
+        session = requests.Session()
+        session.headers.update(self._headers)
+        return session
+
+    def _get_credentials(self) -> dict[str, str]:
+        try:
+            response = self._http.get(self._endpoint, timeout=3)
+            response.raise_for_status()
+        except requests.RequestException as exc:
+            raise BRSError(
+                f"Failed to retrieve ECS credentials from {self._endpoint}"
+            ) from exc
+
+        credentials = response.json()
+        required = {
+            "AccessKeyId",
+            "SecretAccessKey",
+            "SessionToken",
+            "Expiration",
+        }
+        if not required.issubset(credentials):
+            raise BRSError(f"Incomplete credentials received: {credentials}")
+        return {
+            "access_key": credentials.get("AccessKeyId"),
+            "secret_key": credentials.get("SecretAccessKey"),
+            "token": credentials.get("SessionToken"),
+            "expiry_time": credentials.get("Expiration"),  # already ISO8601
+        }
+
+    @staticmethod
+    def get_identity() -> dict[str, str]:
+        """Returns metadata about ECS.
+
+        Returns
+        -------
+        dict[str, str]
+            Dict containing metadata about ECS.
+        """
+
+        return {"method": "ecs", "source": "ecs-container-metadata"}

boto3_refresh_session/exceptions.py

@@ -0,0 +1,38 @@
+class BRSError(Exception):
+    """The base exception for boto3-refresh-session.
+
+    Parameters
+    ----------
+    message : str, optional
+        The message to raise.
+    """
+
+    def __init__(self, message: str | None = None):
+        self.message = "" if message is None else message
+        super().__init__(self.message)
+
+    def __str__(self) -> str:
+        return self.message
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({repr(self.message)})"
+
+
+class BRSWarning(UserWarning):
+    """The base warning for boto3-refresh-session.
+
+    Parameters
+    ----------
+    message : str, optional
+        The message to raise.
+    """
+
+    def __init__(self, message: str | None = None):
+        self.message = "" if message is None else message
+        super().__init__(self.message)
+
+    def __str__(self) -> str:
+        return self.message
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({repr(self.message)})"

boto3_refresh_session/session.py

@@ -1,54 +1,21 @@
 from __future__ import annotations
 
-__doc__ = """
-boto3_refresh_session.session
-=============================
-
-This module provides the main interface for constructing refreshable boto3 sessions.
-
-The ``RefreshableSession`` class serves as a factory that dynamically selects the appropriate 
-credential refresh strategy based on the ``method`` parameter, e.g., ``sts``.
-
-Users can interact with AWS services just like they would with a normal :class:`boto3.session.Session`, 
-with the added benefit of automatic credential refreshing.
-
-Examples
---------
->>> from boto3_refresh_session import RefreshableSession
->>> session = RefreshableSession(
-...     assume_role_kwargs={"RoleArn": "...", "RoleSessionName": "..."},
-...     region_name="us-east-1"
-... )
->>> s3 = session.client("s3")
->>> s3.list_buckets()
-
-.. seealso::
-    :class:`boto3_refresh_session.sts.STSRefreshableSession`
-
-Factory interface
------------------
-.. autosummary::
-   :toctree: generated/
-   :nosignatures:
-
-   RefreshableSession
-"""
-
 __all__ = ["RefreshableSession"]
 
 from abc import ABC, abstractmethod
 from typing import Any, Callable, ClassVar, Literal, get_args
-from warnings import warn
 
 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"]
-RefreshMethod = Literal["sts-assume-role"]
+Method = Literal["sts", "ecs"]
+RefreshMethod = Literal["sts-assume-role", "ecs-container-metadata"]
 
 
 class BaseRefreshableSession(ABC, Session):
@@ -77,7 +44,9 @@ def __init_subclass__(cls, method: Method):
 
         # guarantees that methods are unique
         if method in BaseRefreshableSession.registry:
-            warn(f"Method '{method}' is already registered. Overwriting.")
+            BRSWarning(
+                f"Method {repr(method)} is already registered. Overwriting."
+            )
 
         BaseRefreshableSession.registry[method] = cls
 
@@ -108,6 +77,34 @@ def _refresh_using(
                 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.
+        """
+
+        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,
+        }
+
+    @property
+    def credentials(self) -> dict[str, str]:
+        """The current temporary AWS security credentials."""
+
+        return self.refreshable_credentials()
+
 
 class RefreshableSession:
     """Factory class for constructing refreshable boto3 sessions using various authentication
@@ -134,11 +131,18 @@ class RefreshableSession:
     See Also
     --------
     boto3_refresh_session.sts.STSRefreshableSession
+    boto3_refresh_session.ecs.ECSRefreshableSession
     """
 
     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)}."
+            )
+
         obj = BaseRefreshableSession.registry[method]
         return obj(**kwargs)
 

boto3_refresh_session/sts.py

@@ -1,49 +1,10 @@
 from __future__ import annotations
 
-__doc__ = """
-boto3_refresh_session.sts
-=========================
-
-Implements the STS-based credential refresh strategy for use with 
-:class:`boto3_refresh_session.session.RefreshableSession`.
-
-This module defines the :class:`STSRefreshableSession` class, which uses 
-IAM role assumption via STS to automatically refresh temporary credentials 
-in the background.
-
-.. versionadded:: 1.1.0
-
-Examples
---------
->>> from boto3_refresh_session import RefreshableSession
->>> session = RefreshableSession(
-...     method="sts",
-...     assume_role_kwargs={
-...         "RoleArn": "arn:aws:iam::123456789012:role/MyRole",
-...         "RoleSessionName": "my-session"
-...     },
-...     region_name="us-east-1"
-... )
->>> s3 = session.client("s3")
->>> s3.list_buckets()
-
-.. seealso::
-    :class:`boto3_refresh_session.session.RefreshableSession`
-
-STS
----    
-
-.. autosummary::
-   :toctree: generated/
-   :nosignatures:
-
-   STSRefreshableSession
-"""
 __all__ = ["STSRefreshableSession"]
 
 from typing import Any
-from warnings import warn
 
+from .exceptions import BRSWarning
 from .session import BaseRefreshableSession
 
 
@@ -73,7 +34,7 @@ class STSRefreshableSession(BaseRefreshableSession, method="sts"):
     def __init__(
         self,
         assume_role_kwargs: dict,
-        defer_refresh: bool = None,
+        defer_refresh: bool | None = None,
         sts_client_kwargs: dict | None = None,
         **kwargs,
     ):
@@ -84,8 +45,8 @@ def __init__(
         if sts_client_kwargs is not None:
             # overwriting 'service_name' in case it appears in sts_client_kwargs
             if "service_name" in sts_client_kwargs:
-                warn(
-                    "The sts_client_kwargs parameter cannot contain values for service_name. Reverting to service_name = 'sts'."
+                BRSWarning(
+                    "'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(

doc/conf.py

@@ -82,10 +82,10 @@
     "members": True,
     "member-order": "bysource",
     "exclude-members": "__init__,__new__",
+    "inherited-members": True,
 }
-autodoc_typehints = "none"
-autodoc_preserve_defaults = False
-autodoc_class_signature = "separated"
+autodoc_typehints = "signature"
+autodoc_inherit_docstrings = True
 
 # numpydoc config
 numpydoc_show_class_members = False