amazon-ads-mcp 0.2.7__py3-none-any.whl

amazon_ads_mcp/middleware/caching.py

@@ -0,0 +1,177 @@
+"""Response caching middleware for Amazon Ads MCP.
+
+This module provides a security-aware caching configuration that prevents
+cross-account data leakage in multi-tenant scenarios.
+
+Security Considerations
+-----------------------
+Amazon Ads MCP operates in a multi-tenant context where:
+- Different profiles have different data access
+- Region affects API endpoints and responses
+- Account ID determines data isolation
+
+The default FastMCP cache key (method + arguments) is UNSAFE because:
+- Profile context is implicit (via Amazon-Advertising-API-Scope header)
+- Results vary by active profile, but cache key doesn't include it
+- OpenBridge identities add another dimension of isolation
+
+Safe Caching Strategy
+---------------------
+1. WHITELIST ONLY - explicit list of safe-to-cache tools
+2. STATIC DATA ONLY - tools that return the same data regardless of context
+3. NO API CALLS - only cache server-local metadata
+
+Examples
+--------
+.. code-block:: python
+
+    from amazon_ads_mcp.middleware.caching import create_caching_middleware
+
+    middleware = create_caching_middleware()
+    server.add_middleware(middleware)
+"""
+
+import logging
+from typing import Optional, Set
+
+from fastmcp.server.middleware.caching import (
+    CallToolSettings,
+    ListPromptsSettings,
+    ListResourcesSettings,
+    ListToolsSettings,
+    ResponseCachingMiddleware,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# Tools that are SAFE to cache (static data, no profile dependency)
+SAFE_TO_CACHE_TOOLS: Set[str] = {
+    # Region configuration - static server data
+    "list_regions",
+    # Downloads - local filesystem, not API-dependent
+    "list_downloads",
+}
+
+# Tools that MUST NOT be cached (profile-dependent, write operations, or dynamic)
+# This is a documentation list - actual enforcement is via whitelist above
+NEVER_CACHE_TOOLS: Set[str] = {
+    # Profile/Identity management - state changes
+    "set_active_profile",
+    "get_active_profile",  # Depends on server state
+    "clear_active_profile",
+    "set_active_identity",
+    "get_active_identity",  # Depends on server state
+    "list_identities",  # Depends on auth provider
+    # Region management - state changes
+    "set_region",
+    "get_region",  # Depends on server state
+    "get_routing_state",  # Depends on current routing config
+    # OAuth operations - security sensitive
+    "start_oauth_flow",
+    "check_oauth_status",
+    "refresh_oauth_token",
+    "clear_oauth_tokens",
+    # Download operations - side effects
+    "download_export",
+    # Sampling - dynamic operations
+    "test_sampling",
+    # ALL OpenAPI-generated tools - profile-dependent API responses
+    # These are excluded by whitelist (not in SAFE_TO_CACHE_TOOLS)
+}
+
+# TTL values in seconds
+STATIC_DATA_TTL = 3600  # 1 hour for truly static data
+LIST_METADATA_TTL = 60  # 1 minute for tool/resource/prompt lists
+
+
+def create_caching_middleware(
+    enabled: bool = True,
+    static_ttl: int = STATIC_DATA_TTL,
+    list_ttl: int = LIST_METADATA_TTL,
+    additional_safe_tools: Optional[Set[str]] = None,
+) -> ResponseCachingMiddleware:
+    """Create a security-aware caching middleware.
+
+    This middleware implements a conservative whitelist approach:
+    - Only explicitly listed tools are cached
+    - All OpenAPI-generated tools are excluded (profile-dependent)
+    - Write operations are never cached
+
+    :param enabled: Whether caching is enabled globally
+    :param static_ttl: TTL for static data (e.g., list_regions)
+    :param list_ttl: TTL for list operations (tools, resources, prompts)
+    :param additional_safe_tools: Additional tool names safe to cache
+    :return: Configured ResponseCachingMiddleware
+
+    Example
+    -------
+    .. code-block:: python
+
+        middleware = create_caching_middleware(
+            static_ttl=1800,  # 30 minutes
+            additional_safe_tools={"my_static_tool"}
+        )
+        server.add_middleware(middleware)
+    """
+    safe_tools = SAFE_TO_CACHE_TOOLS.copy()
+    if additional_safe_tools:
+        safe_tools.update(additional_safe_tools)
+
+    logger.info(
+        f"Creating caching middleware with {len(safe_tools)} safe tools: {safe_tools}"
+    )
+
+    return ResponseCachingMiddleware(
+        # Tool call caching - WHITELIST ONLY
+        call_tool_settings=CallToolSettings(
+            enabled=enabled,
+            ttl=static_ttl,
+            included_tools=list(safe_tools),  # Only these tools are cached
+        ),
+        # List operations - safe to cache (server metadata)
+        list_tools_settings=ListToolsSettings(
+            enabled=enabled,
+            ttl=list_ttl,
+        ),
+        list_resources_settings=ListResourcesSettings(
+            enabled=enabled,
+            ttl=list_ttl,
+        ),
+        list_prompts_settings=ListPromptsSettings(
+            enabled=enabled,
+            ttl=list_ttl,
+        ),
+        # Resource reads - DISABLED (may be profile-dependent)
+        # read_resource_settings=ReadResourceSettings(enabled=False),
+        # Prompt gets - DISABLED (may be profile-dependent)
+        # get_prompt_settings=GetPromptSettings(enabled=False),
+    )
+
+
+# Future enhancement: Custom cache key middleware
+# This would allow caching OpenAPI tools safely by including
+# profile_id/region/account_id in the cache key
+#
+# class ContextAwareCacheKeyMiddleware(Middleware):
+#     """Middleware that injects routing context into cache keys.
+#
+#     This middleware captures the current profile_id, region, and
+#     account_id and stores them in context state for use by a
+#     custom caching implementation.
+#     """
+#
+#     async def on_call_tool(self, context: MiddlewareContext, call_next):
+#         # Get current routing context
+#         from ..utils.http_client import get_routing_state
+#         routing = get_routing_state()
+#
+#         # Store in context for cache key generation
+#         if context.fastmcp_context:
+#             context.fastmcp_context.set_state("cache_context", {
+#                 "region": routing.get("region"),
+#                 "profile_id": routing.get("profile_id"),
+#                 "account_id": routing.get("account_id"),
+#             })
+#
+#         return await call_next(context)

amazon_ads_mcp/middleware/oauth.py

@@ -0,0 +1,175 @@
+"""OAuth middleware for automatic token injection."""
+
+import logging
+from datetime import datetime, timedelta, timezone
+from typing import Optional
+
+import httpx
+from fastmcp.server.middleware import Middleware, MiddlewareContext
+
+from ..tools.oauth import OAuthTokens
+from ..utils.region_config import RegionConfig
+
+logger = logging.getLogger(__name__)
+
+
+class OAuthTokenMiddleware(Middleware):
+    """
+    Middleware that automatically injects OAuth tokens into API calls.
+
+    This middleware:
+    1. Checks for stored OAuth tokens in the context state
+    2. Refreshes expired access tokens automatically
+    3. Injects tokens into the authentication flow
+
+    Note: This middleware uses the AuthManager public API:
+    - get_active_identity() to check current identity
+    - set_active_identity() to switch to OAuth identity
+    - Stores tokens in context state for providers to access
+    """
+
+    def __init__(self, client_id: str, client_secret: str, region: str = "na"):
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.region = region
+
+    async def refresh_token(self, refresh_token: str) -> Optional[dict]:
+        """Refresh an expired access token."""
+        token_url = RegionConfig.get_oauth_endpoint(self.region)
+        token_data = {
+            "grant_type": "refresh_token",
+            "refresh_token": refresh_token,
+            "client_id": self.client_id,
+            "client_secret": self.client_secret,
+        }
+
+        try:
+            # Use explicit timeout for OAuth token refresh
+            timeout = httpx.Timeout(connect=10.0, read=30.0, write=10.0, pool=10.0)
+            async with httpx.AsyncClient(timeout=timeout) as client:
+                response = await client.post(token_url, data=token_data)
+
+            if response.status_code == 200:
+                return response.json()
+            else:
+                logger.error(
+                    f"Failed to refresh token: {response.status_code} - {response.text}"
+                )
+                return None
+        except Exception as e:
+            logger.error(f"Error refreshing token: {e}")
+            return None
+
+    async def on_call_tool(self, context: MiddlewareContext, call_next):
+        """
+        Intercept tool calls to inject OAuth tokens if available.
+        """
+        # Skip OAuth tools themselves to avoid recursion
+        if context.message and hasattr(context.message, "name"):
+            tool_name = context.message.name
+            if tool_name and "oauth" in tool_name.lower():
+                return await call_next(context)
+
+        # Check for OAuth tokens in state
+        if context.fastmcp_context:
+            try:
+                tokens_data = await context.fastmcp_context.get_state("oauth_tokens")
+
+                if tokens_data:
+                    tokens = OAuthTokens(**tokens_data)
+
+                    # Check if token needs refresh
+                    if tokens.is_expired and tokens.refresh_token:
+                        logger.info("OAuth access token expired, refreshing...")
+
+                        token_response = await self.refresh_token(tokens.refresh_token)
+                        if token_response:
+                            # Update tokens
+                            tokens.access_token = token_response["access_token"]
+                            tokens.expires_in = token_response.get("expires_in", 3600)
+                            tokens.obtained_at = datetime.now(timezone.utc)
+
+                            if "refresh_token" in token_response:
+                                tokens.refresh_token = token_response["refresh_token"]
+
+                            # Store updated tokens
+                            await context.fastmcp_context.set_state(
+                                "oauth_tokens", tokens.model_dump()
+                            )
+                            logger.info("OAuth access token refreshed successfully")
+
+                    # If auth manager exists, store tokens through unified token store
+                    if hasattr(context.fastmcp_context, "auth_manager"):
+                        auth_manager = context.fastmcp_context.auth_manager
+
+                        # Store tokens in unified token store
+                        if hasattr(auth_manager, "set_token"):
+                            from ..auth.token_store import TokenKind
+
+                            # Store access token
+                            expires_at = tokens.obtained_at + timedelta(
+                                seconds=tokens.expires_in
+                            )
+                            await auth_manager.set_token(
+                                provider_type="oauth",
+                                identity_id="oauth",
+                                token_kind=TokenKind.ACCESS,
+                                token=tokens.access_token,
+                                expires_at=expires_at,
+                                metadata={"token_type": "Bearer"},
+                            )
+
+                            # Store refresh token
+                            await auth_manager.set_token(
+                                provider_type="oauth",
+                                identity_id="oauth",
+                                token_kind=TokenKind.REFRESH,
+                                token=tokens.refresh_token,
+                                expires_at=datetime.now(timezone.utc)
+                                + timedelta(days=365),  # Long-lived
+                                metadata={},
+                            )
+                            logger.debug("Stored OAuth tokens in unified token store")
+
+                        # Check current active identity
+                        active_identity = auth_manager.get_active_identity()
+
+                        # If not using OAuth identity, try to switch
+                        if not active_identity or active_identity.id != "oauth":
+                            try:
+                                # Try to set OAuth as active identity
+                                # This assumes OAuth provider is configured or identity exists
+                                await auth_manager.set_active_identity("oauth")
+                                logger.info("Switched to OAuth authentication identity")
+                            except Exception as e:
+                                # OAuth identity doesn't exist or provider not configured for it
+                                logger.debug(f"Could not switch to OAuth identity: {e}")
+                    else:
+                        # Fallback: Store tokens in context for backward compatibility
+                        await context.fastmcp_context.set_state(
+                            "current_access_token", tokens.access_token
+                        )
+                        await context.fastmcp_context.set_state(
+                            "current_refresh_token", tokens.refresh_token
+                        )
+
+            except Exception as e:
+                logger.debug(f"OAuth middleware check: {e}")
+                # Continue without OAuth tokens
+
+        # Continue with the tool call
+        return await call_next(context)
+
+
+def create_oauth_middleware():
+    """Create OAuth middleware instance with settings."""
+    from ..config.settings import settings
+
+    if not settings.oauth_client_id or not settings.oauth_client_secret:
+        logger.warning("OAuth client credentials not configured")
+        return None
+
+    return OAuthTokenMiddleware(
+        client_id=settings.oauth_client_id,
+        client_secret=settings.oauth_client_secret,
+    )

amazon_ads_mcp/middleware/sampling.py

@@ -0,0 +1,112 @@
+"""Middleware to attach server-side sampling handler to request context."""
+
+import logging
+from typing import Any, Callable
+
+from fastmcp import Context
+
+logger = logging.getLogger(__name__)
+
+
+def create_sampling_middleware(sampling_handler: Any = None) -> Callable:
+    """
+    Create middleware that attaches the server's sampling handler to each request context.
+
+    This allows sample_with_fallback() to discover the handler when the client
+    doesn't support sampling.
+
+    Args:
+        sampling_handler: The server's sampling handler instance
+
+    Returns:
+        Middleware function that can be added to the server
+    """
+
+    async def sampling_middleware(request: Any, handler: Callable) -> Any:
+        """
+        Middleware that attaches sampling handler to the request context.
+
+        Args:
+            request: The incoming request
+            handler: The next handler in the chain
+
+        Returns:
+            Response from the handler chain
+        """
+        # Try to get the context from the request
+        # FastMCP typically stores context in the request or uses a context var
+        try:
+            # Method 1: Check if there's a context attribute on the request
+            if hasattr(request, "context"):
+                ctx = request.context
+                if isinstance(ctx, Context):
+                    # Use the wrapper to provide sampling
+                    from ..utils.sampling_wrapper import get_sampling_wrapper
+
+                    wrapper = get_sampling_wrapper()
+                    if wrapper.has_handler():
+                        # Try to use public API if available, otherwise skip
+                        if hasattr(ctx, "set_sampling_handler"):
+                            ctx.set_sampling_handler(wrapper)
+                        else:
+                            # Avoid setting private attributes
+                            logger.debug(
+                                "Skipping sampling attachment - no public API available"
+                            )
+                    logger.debug("Processed sampling handler for request context")
+
+            # Method 2: Check for FastMCP context in request state
+            elif hasattr(request, "state") and hasattr(
+                request.state, "fastmcp_context"
+            ):
+                ctx = request.state.fastmcp_context
+                if isinstance(ctx, Context):
+                    from ..utils.sampling_wrapper import get_sampling_wrapper
+
+                    wrapper = get_sampling_wrapper()
+                    if wrapper.has_handler():
+                        # Try to use public API if available, otherwise skip
+                        if hasattr(ctx, "set_sampling_handler"):
+                            ctx.set_sampling_handler(wrapper)
+                        else:
+                            # Avoid setting private attributes
+                            logger.debug(
+                                "Skipping sampling attachment - no public API available"
+                            )
+                    logger.debug("Processed sampling handler for FastMCP context")
+
+            # Method 3: Skip private contextvar usage
+            else:
+                # Avoid using private _current_context API
+                logger.debug("Skipping contextvar method to avoid private API usage")
+
+        except Exception as e:
+            logger.debug(f"Could not attach sampling handler to context: {e}")
+
+        # Continue with the request
+        return await handler(request)
+
+    return sampling_middleware
+
+
+def attach_sampling_to_context(ctx: Context) -> None:
+    """
+    Helper function to directly attach sampling handler to a context.
+
+    This can be called from tool handlers or other places where we have
+    direct access to the context.
+
+    Args:
+        ctx: The FastMCP context
+    """
+    from ..utils.sampling_wrapper import get_sampling_wrapper
+
+    wrapper = get_sampling_wrapper()
+    if wrapper.has_handler():
+        # Try to use public API if available, otherwise skip
+        if hasattr(ctx, "set_sampling_handler"):
+            ctx.set_sampling_handler(wrapper)
+            logger.debug("Sampling handler attached to context")
+        else:
+            # Avoid setting private attributes
+            logger.debug("Skipping sampling attachment - no public API available")