""" OAuth2 client-credentials token provider for the Amazon Shuttle API. Amazon issues 60-minute JWT access tokens from a Cognito endpoint. Tokens must be cached and reused: requesting one per API call gets the client rate-limited. This provider fetches lazily, hands out the cached token until it is close to expiry (Amazon recommends refreshing at ~50 minutes), and serialises concurrent refreshes so a burst of webhooks triggers one fetch. """ import asyncio import base64 import logging import time from typing import Callable, Optional import aiohttp logger = logging.getLogger(__name__) # Refresh this many seconds before the token actually expires. DEFAULT_REFRESH_MARGIN_SECONDS = 600 DEFAULT_TOKEN_LIFETIME_SECONDS = 3600 class AmazonAuthError(Exception): """Raised when an access token could not be obtained from Cognito. retryable=True marks transport failures and 5xx/429 responses, where a later attempt may succeed. Anything else (400 malformed request, 401 invalid_client) is a configuration problem that retrying will not fix. """ def __init__(self, message: str, status: Optional[int] = None, retryable: bool = False): super().__init__(message) self.status = status self.retryable = retryable class CognitoTokenProvider: def __init__( self, token_url: str, client_id: str, client_secret: str, scope: str, *, refresh_margin_seconds: int = DEFAULT_REFRESH_MARGIN_SECONDS, timeout_seconds: float = 10.0, clock: Callable[[], float] = time.monotonic, ): if not all((token_url, client_id, client_secret, scope)): raise ValueError( "CognitoTokenProvider requires token_url, client_id, client_secret and scope" ) self._token_url = token_url self._scope = scope self._basic_auth = base64.b64encode( f"{client_id}:{client_secret}".encode("utf-8") ).decode("ascii") self._refresh_margin = refresh_margin_seconds self._timeout = aiohttp.ClientTimeout(total=timeout_seconds) self._clock = clock self._access_token: Optional[str] = None self._refresh_at: float = 0.0 self._expires_at: float = 0.0 self._lock: Optional[asyncio.Lock] = None @classmethod def from_config(cls, config) -> "CognitoTokenProvider": return cls( token_url=config.amazon_cognito_token_url, client_id=config.amazon_client_id, client_secret=config.amazon_client_secret, scope=config.amazon_oauth_scope, ) @property def has_token(self) -> bool: return self._access_token is not None async def get_token(self) -> str: """Return a cached access token, fetching or refreshing when needed.""" if self._is_fresh(): return self._access_token async with self._get_lock(): if self._is_fresh(): return self._access_token try: await self._fetch_token() except AmazonAuthError: # A failed early refresh should not take down requests while # the current token is still inside its real lifetime. if self._is_still_valid(): logger.warning( "Amazon token refresh failed; reusing current token for up to %ds", int(self._expires_at - self._clock()), ) return self._access_token raise return self._access_token def invalidate(self) -> None: """Drop the cached token so the next get_token() fetches a new one. Call this after Amazon rejects a request with 401 so the retry uses a fresh token; do not loop on it if the second attempt also fails. """ if self._access_token is not None: logger.info("Invalidated cached Amazon access token") self._access_token = None self._refresh_at = 0.0 self._expires_at = 0.0 def _is_fresh(self) -> bool: return self._access_token is not None and self._clock() < self._refresh_at def _is_still_valid(self) -> bool: return self._access_token is not None and self._clock() < self._expires_at def _get_lock(self) -> asyncio.Lock: # Created lazily so the lock binds to the running loop; a lock built # in __init__ can be tied to a different loop on older Pythons. if self._lock is None: self._lock = asyncio.Lock() return self._lock async def _fetch_token(self) -> None: headers = { "Authorization": f"Basic {self._basic_auth}", "Content-Type": "application/x-www-form-urlencoded", } data = {"grant_type": "client_credentials", "scope": self._scope} try: async with aiohttp.ClientSession(timeout=self._timeout) as session: async with session.post( self._token_url, data=data, headers=headers ) as response: if response.status != 200: body = await response.text() retryable = response.status >= 500 or response.status == 429 logger.error( "Amazon Cognito token request failed: HTTP %s %s", response.status, body[:300], ) raise AmazonAuthError( f"Cognito token request failed with HTTP {response.status}", status=response.status, retryable=retryable, ) payload = await response.json() except (aiohttp.ClientError, asyncio.TimeoutError) as e: logger.error("Amazon Cognito token request failed: %s", e) raise AmazonAuthError( f"Cognito token request failed: {e}", retryable=True ) from e token = payload.get("access_token") if isinstance(payload, dict) else None if not token: raise AmazonAuthError("Cognito response did not include access_token") try: expires_in = int(payload.get("expires_in", DEFAULT_TOKEN_LIFETIME_SECONDS)) except (TypeError, ValueError): expires_in = DEFAULT_TOKEN_LIFETIME_SECONDS # Never let the margin swallow the whole lifetime: a short-lived token # is refreshed at half-life instead of on every call. margin = min(self._refresh_margin, expires_in // 2) now = self._clock() self._access_token = token self._expires_at = now + expires_in self._refresh_at = now + expires_in - margin logger.info( "Fetched Amazon access token (expires_in=%ss, refresh_in=%ss)", expires_in, expires_in - margin, )