msgflux 0.1.0__py3-none-any.whl

msgflux/protocols/mcp/auth/bearer.py

@@ -0,0 +1,97 @@
+"""Bearer token authentication for MCP."""
+
+from datetime import datetime
+from typing import Awaitable, Callable, Dict, Optional
+
+from .base import BaseAuth
+
+
+class BearerTokenAuth(BaseAuth):
+    """Bearer token authentication (JWT, etc).
+
+    This auth provider adds an Authorization header with a Bearer token.
+    Supports automatic token refresh via callback.
+
+    Example:
+        ```python
+        auth = BearerTokenAuth(token="your-jwt-token")
+
+        # With auto-refresh
+        async def get_new_token():
+            # Your token refresh logic
+            return "new-token"
+
+        auth = BearerTokenAuth(
+            token="initial-token",
+            expires_in=3600,
+            refresh_callback=get_new_token
+        )
+        ```
+    """
+
+    def __init__(
+        self,
+        token: str,
+        expires_in: Optional[int] = None,
+        refresh_callback: Optional[Callable[[], Awaitable[str]]] = None,
+        token_type: str = "Bearer",  # noqa: S107,
+    ):
+        """Initialize bearer token authentication.
+
+        Args:
+            token: The bearer token.
+            expires_in: Token expiration time in seconds (optional).
+            refresh_callback: Async callback to get a new token when expired.
+            token_type: Type of token (default: "Bearer").
+        """
+        super().__init__()
+        self._token = token
+        self._token_type = token_type
+        self._refresh_callback = refresh_callback
+
+        if expires_in:
+            self.set_expiration(expires_in)
+
+    def apply_auth(self, headers: Dict[str, str]) -> Dict[str, str]:
+        """Apply bearer token to Authorization header.
+
+        Args:
+            headers: Existing request headers.
+
+        Returns:
+            Headers with Authorization added.
+        """
+        headers = headers.copy()
+        headers["Authorization"] = f"{self._token_type} {self._token}"
+        return headers
+
+    async def refresh(self) -> bool:
+        """Refresh the bearer token using the callback.
+
+        Returns:
+            True if refresh succeeded, False otherwise.
+        """
+        if not self._refresh_callback:
+            return False
+
+        try:
+            new_token = await self._refresh_callback()
+            if new_token:
+                self._token = new_token
+                self._last_refresh = datetime.now()  # noqa: DTZ005
+                return True
+        except Exception:  # noqa: S110
+            pass
+
+        return False
+
+    def update_token(self, token: str, expires_in: Optional[int] = None) -> None:
+        """Update the bearer token.
+
+        Args:
+            token: New bearer token.
+            expires_in: New expiration time in seconds (optional).
+        """
+        self._token = token
+        if expires_in:
+            self.set_expiration(expires_in)

msgflux/protocols/mcp/auth/custom.py

@@ -0,0 +1,76 @@
+"""Custom header authentication for MCP."""
+
+from typing import Callable, Dict, Optional
+
+from .base import BaseAuth
+
+
+class CustomHeaderAuth(BaseAuth):
+    """Custom header-based authentication.
+
+    Allows arbitrary headers to be added for authentication.
+    Useful for proprietary auth schemes or multiple headers.
+
+    Example:
+        ```python
+        # Static headers
+        auth = CustomHeaderAuth({
+            "X-Custom-Auth": "secret-value",
+            "X-Request-ID": "unique-id"
+        })
+
+        # Dynamic headers via callback
+        def get_headers():
+            return {
+                "X-Signature": compute_signature(),
+                "X-Timestamp": str(int(time.time()))
+            }
+
+        auth = CustomHeaderAuth(headers_callback=get_headers)
+        ```
+    """
+
+    def __init__(
+        self,
+        headers: Optional[Dict[str, str]] = None,
+        headers_callback: Optional[Callable[[], Dict[str, str]]] = None,
+    ):
+        """Initialize custom header authentication.
+
+        Args:
+            headers: Static headers to add (optional).
+            headers_callback: Callback that returns headers dynamically (optional).
+                            Takes precedence over static headers if both provided.
+        """
+        super().__init__()
+        self._headers = headers or {}
+        self._headers_callback = headers_callback
+
+    def apply_auth(self, headers: Dict[str, str]) -> Dict[str, str]:
+        """Apply custom headers to request.
+
+        Args:
+            headers: Existing request headers.
+
+        Returns:
+            Headers with custom auth headers added.
+        """
+        headers = headers.copy()
+
+        # Use callback if provided, otherwise use static headers
+        if self._headers_callback:
+            auth_headers = self._headers_callback()
+        else:
+            auth_headers = self._headers
+
+        # Merge auth headers
+        headers.update(auth_headers)
+        return headers
+
+    def update_headers(self, headers: Dict[str, str]) -> None:
+        """Update static headers.
+
+        Args:
+            headers: New headers to use.
+        """
+        self._headers = headers

msgflux/protocols/mcp/auth/oauth2.py

@@ -0,0 +1,133 @@
+"""OAuth2 authentication for MCP."""
+
+from datetime import datetime
+from typing import Any, Awaitable, Callable, Dict, Optional
+
+from msgflux.protocols.mcp.auth.base import BaseAuth
+
+
+class OAuth2Auth(BaseAuth):
+    """OAuth2 authentication with automatic token refresh.
+
+    Manages OAuth2 access tokens and refresh tokens, automatically
+    refreshing when expired.
+
+    Example:
+        ```python
+        async def refresh_token_callback(refresh_token: str) -> dict:
+            # Call your OAuth2 token endpoint
+            response = await http_client.post(
+                "https://auth.example.com/oauth/token",
+                data={
+                    "grant_type": "refresh_token",
+                    "refresh_token": refresh_token,
+                    "client_id": "your-client-id",
+                }
+            )
+            return {
+                "access_token": response["access_token"],
+                "refresh_token": response.get("refresh_token"),
+                "expires_in": response["expires_in"],
+            }
+
+        auth = OAuth2Auth(
+            access_token="initial-token",
+            refresh_token="refresh-token",
+            expires_in=3600,
+            refresh_callback=refresh_token_callback
+        )
+        ```
+    """
+
+    def __init__(
+        self,
+        access_token: str,
+        refresh_token: Optional[str] = None,
+        expires_in: Optional[int] = None,
+        refresh_callback: Optional[Callable[[str], Awaitable[Dict[str, Any]]]] = None,
+        token_type: str = "Bearer",  # noqa: S107,
+    ):
+        """Initialize OAuth2 authentication.
+
+        Args:
+            access_token: The OAuth2 access token.
+            refresh_token: The refresh token (required for auto-refresh).
+            expires_in: Token expiration time in seconds.
+            refresh_callback: Async callback that takes refresh_token and returns
+                            dict with 'access_token', 'refresh_token' (optional),
+                            and 'expires_in'.
+            token_type: Type of token (default: "Bearer").
+        """
+        super().__init__()
+        self._access_token = access_token
+        self._refresh_token = refresh_token
+        self._token_type = token_type
+        self._refresh_callback = refresh_callback
+
+        if expires_in:
+            self.set_expiration(expires_in)
+
+    def apply_auth(self, headers: Dict[str, str]) -> Dict[str, str]:
+        """Apply OAuth2 access token to Authorization header.
+
+        Args:
+            headers: Existing request headers.
+
+        Returns:
+            Headers with Authorization added.
+        """
+        headers = headers.copy()
+        headers["Authorization"] = f"{self._token_type} {self._access_token}"
+        return headers
+
+    async def refresh(self) -> bool:
+        """Refresh the OAuth2 tokens using the refresh token.
+
+        Returns:
+            True if refresh succeeded, False otherwise.
+        """
+        if not self._refresh_token or not self._refresh_callback:
+            return False
+
+        try:
+            # Call refresh callback with current refresh token
+            result = await self._refresh_callback(self._refresh_token)
+
+            # Update tokens
+            if "access_token" in result:
+                self._access_token = result["access_token"]
+
+                # Update refresh token if provided
+                if "refresh_token" in result:
+                    self._refresh_token = result["refresh_token"]
+
+                # Update expiration
+                if "expires_in" in result:
+                    self.set_expiration(result["expires_in"])
+
+                self._last_refresh = datetime.now()  # noqa: DTZ005
+                return True
+
+        except Exception:  # noqa: S110
+            pass
+
+        return False
+
+    def update_tokens(
+        self,
+        access_token: str,
+        refresh_token: Optional[str] = None,
+        expires_in: Optional[int] = None,
+    ) -> None:
+        """Update OAuth2 tokens.
+
+        Args:
+            access_token: New access token.
+            refresh_token: New refresh token (optional).
+            expires_in: New expiration time in seconds (optional).
+        """
+        self._access_token = access_token
+        if refresh_token:
+            self._refresh_token = refresh_token
+        if expires_in:
+            self.set_expiration(expires_in)