nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/authentication/http_basic_auth/register.py

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nat.builder.builder import Builder
+from nat.cli.register_workflow import register_auth_provider
+from nat.data_models.authentication import AuthProviderBaseConfig
+
+
+class HTTPBasicAuthProviderConfig(AuthProviderBaseConfig, name="http_basic"):
+    pass
+
+
+@register_auth_provider(config_type=HTTPBasicAuthProviderConfig)
+async def http_basic_auth_provider(config: HTTPBasicAuthProviderConfig, builder: Builder):
+
+    from nat.authentication.http_basic_auth.http_basic_auth_provider import HTTPBasicAuthProvider
+
+    yield HTTPBasicAuthProvider(config)

nat/authentication/interfaces.py

@@ -0,0 +1,96 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typing
+from abc import ABC
+from abc import abstractmethod
+
+from nat.data_models.authentication import AuthenticatedContext
+from nat.data_models.authentication import AuthFlowType
+from nat.data_models.authentication import AuthProviderBaseConfig
+from nat.data_models.authentication import AuthProviderBaseConfigT
+from nat.data_models.authentication import AuthResult
+
+AUTHORIZATION_HEADER = "Authorization"
+
+
+class AuthProviderBase(typing.Generic[AuthProviderBaseConfigT], ABC):
+    """
+    Base class for authenticating to API services.
+    This class provides an interface for authenticating to API services.
+    """
+
+    def __init__(self, config: AuthProviderBaseConfigT):
+        """
+        Initialize the AuthProviderBase with the given configuration.
+
+        Args:
+            config (AuthProviderBaseConfig): Configuration items for authentication.
+        """
+        self._config = config
+
+    @property
+    def config(self) -> AuthProviderBaseConfigT:
+        """
+        Returns the auth provider configuration object.
+
+        Returns
+        -------
+        AuthProviderBaseConfigT
+            The auth provider configuration object.
+        """
+        return self._config
+
+    @abstractmethod
+    async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult:
+        """
+        Perform the authentication process for the client.
+
+        This method handles the necessary steps to authenticate the client with the
+        target API service, which may include obtaining tokens, refreshing credentials,
+        or completing multi-step authentication flows.
+
+        Args:
+            user_id: Optional user identifier for authentication
+            kwargs: Additional authentication parameters for example: http response (typically from a 401)
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        # This method will call the frontend FlowHandlerBase `authenticate` method
+        pass
+
+
+class FlowHandlerBase(ABC):
+    """
+    Handles front-end specific flows for authentication clients.
+
+    Each front end will define a FlowHandler that will implement the authenticate method.
+
+    The `authenticate` method will be stored as the callback in the ContextState.user_auth_callback
+    """
+
+    @abstractmethod
+    async def authenticate(self, config: AuthProviderBaseConfig, method: AuthFlowType) -> AuthenticatedContext:
+        """
+        Perform the authentication process for the client.
+
+        This method handles the necessary steps to authenticate the client with the
+        target API service, which may include obtaining tokens, refreshing credentials,
+        or completing multistep authentication flows.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        pass

nat/authentication/oauth2/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nat/authentication/oauth2/oauth2_auth_code_flow_provider.py

@@ -0,0 +1,140 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from collections.abc import Awaitable
+from collections.abc import Callable
+from datetime import UTC
+from datetime import datetime
+
+import httpx
+from authlib.integrations.httpx_client import OAuth2Client as AuthlibOAuth2Client
+from pydantic import SecretStr
+
+from nat.authentication.interfaces import AuthProviderBase
+from nat.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
+from nat.builder.context import Context
+from nat.data_models.authentication import AuthenticatedContext
+from nat.data_models.authentication import AuthFlowType
+from nat.data_models.authentication import AuthResult
+from nat.data_models.authentication import BearerTokenCred
+
+logger = logging.getLogger(__name__)
+
+
+class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConfig]):
+
+    def __init__(self, config: OAuth2AuthCodeFlowProviderConfig, token_storage=None):
+        super().__init__(config)
+        self._auth_callback = None
+        # Always use token storage - defaults to in-memory if not provided
+        if token_storage is None:
+            from nat.plugins.mcp.auth.token_storage import InMemoryTokenStorage
+            self._token_storage = InMemoryTokenStorage()
+        else:
+            self._token_storage = token_storage
+
+    async def _attempt_token_refresh(self, user_id: str, auth_result: AuthResult) -> AuthResult | None:
+        refresh_token = auth_result.raw.get("refresh_token")
+        if not isinstance(refresh_token, str):
+            return None
+
+        try:
+            with AuthlibOAuth2Client(
+                    client_id=self.config.client_id,
+                    client_secret=self.config.client_secret,
+            ) as client:
+                new_token_data = client.refresh_token(self.config.token_url, refresh_token=refresh_token)
+
+                expires_at_ts = new_token_data.get("expires_at")
+                new_expires_at = datetime.fromtimestamp(expires_at_ts, tz=UTC) if expires_at_ts else None
+
+            new_auth_result = AuthResult(
+                credentials=[BearerTokenCred(token=SecretStr(new_token_data["access_token"]))],
+                token_expires_at=new_expires_at,
+                raw=new_token_data,
+            )
+
+            await self._token_storage.store(user_id, new_auth_result)
+        except httpx.HTTPStatusError:
+            return None
+        except httpx.RequestError:
+            return None
+        except Exception:
+            # On any other failure, we'll fall back to the full auth flow.
+            return None
+
+        return new_auth_result
+
+    def _set_custom_auth_callback(self,
+                                  auth_callback: Callable[[OAuth2AuthCodeFlowProviderConfig, AuthFlowType],
+                                                          Awaitable[AuthenticatedContext]]):
+        self._auth_callback = auth_callback
+
+    async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult:
+        context = Context.get()
+        if user_id is None and hasattr(context, "metadata") and hasattr(
+                context.metadata, "cookies") and context.metadata.cookies is not None:
+            session_id = context.metadata.cookies.get("nat-session", None)
+            if not session_id:
+                raise RuntimeError("Authentication failed. No session ID found. Cannot identify user.")
+
+            user_id = session_id
+
+        if user_id:
+            # Try to retrieve from token storage
+            auth_result = await self._token_storage.retrieve(user_id)
+
+            if auth_result:
+                if not auth_result.is_expired():
+                    return auth_result
+
+                refreshed_auth_result = await self._attempt_token_refresh(user_id, auth_result)
+                if refreshed_auth_result:
+                    return refreshed_auth_result
+
+        # Try getting callback from the context if that's not set, use the default callback
+        try:
+            auth_callback = Context.get().user_auth_callback
+        except RuntimeError:
+            auth_callback = self._auth_callback
+
+        if not auth_callback:
+            raise RuntimeError("Authentication callback not set on Context.")
+
+        try:
+            authenticated_context = await auth_callback(self.config, AuthFlowType.OAUTH2_AUTHORIZATION_CODE)
+        except Exception as e:
+            raise RuntimeError(f"Authentication callback failed: {e}") from e
+
+        headers = authenticated_context.headers or {}
+        auth_header = headers.get("Authorization", "")
+        if not auth_header.startswith("Bearer "):
+            raise RuntimeError("Invalid Authorization header")
+
+        token = auth_header.split(" ")[1]
+
+        # Safely access metadata
+        metadata = authenticated_context.metadata or {}
+        auth_result = AuthResult(
+            credentials=[BearerTokenCred(token=SecretStr(token))],
+            token_expires_at=metadata.get("expires_at"),
+            raw=metadata.get("raw_token") or {},
+        )
+
+        if user_id:
+            await self._token_storage.store(user_id, auth_result)
+
+        return auth_result

nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.py

@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pydantic import Field
+
+from nat.data_models.authentication import AuthProviderBaseConfig
+
+
+class OAuth2AuthCodeFlowProviderConfig(AuthProviderBaseConfig, name="oauth2_auth_code_flow"):
+
+    client_id: str = Field(description="The client ID for OAuth 2.0 authentication.")
+    client_secret: str = Field(description="The secret associated with the client_id.")
+    authorization_url: str = Field(description="The authorization URL for OAuth 2.0 authentication.")
+    token_url: str = Field(description="The token URL for OAuth 2.0 authentication.")
+    token_endpoint_auth_method: str | None = Field(
+        description=("The authentication method for the token endpoint. "
+                     "Usually one of `client_secret_post` or `client_secret_basic`."),
+        default=None)
+    redirect_uri: str = Field(description="The redirect URI for OAuth 2.0 authentication. Must match the registered "
+                              "redirect URI with the OAuth provider.")
+    scopes: list[str] = Field(description="The scopes for OAuth 2.0 authentication.", default_factory=list)
+    use_pkce: bool = Field(default=False,
+                           description="Whether to use PKCE (Proof Key for Code Exchange) in the OAuth 2.0 flow.")
+
+    authorization_kwargs: dict[str, str] | None = Field(description=("Additional keyword arguments for the "
+                                                                     "authorization request."),
+                                                        default=None)

nat/authentication/oauth2/oauth2_resource_server_config.py

@@ -0,0 +1,124 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from urllib.parse import urlparse
+
+from pydantic import Field
+from pydantic import field_validator
+from pydantic import model_validator
+
+from nat.data_models.authentication import AuthProviderBaseConfig
+
+
+class OAuth2ResourceServerConfig(AuthProviderBaseConfig, name="oauth2_resource_server"):
+    """OAuth 2.0 Resource Server authentication configuration.
+
+    Supports:
+      • JWT access tokens via JWKS / OIDC Discovery / issuer fallback
+      • Opaque access tokens via RFC 7662 introspection
+    """
+
+    issuer_url: str = Field(
+        description=("The unique issuer identifier for an authorization server. "
+                     "Required for validation and used to derive the default JWKS URI "
+                     "(<issuer_url>/.well-known/jwks.json) if `jwks_uri` and `discovery_url` are not provided."), )
+    scopes: list[str] = Field(
+        default_factory=list,
+        description="Scopes required by this API. Validation ensures the token grants all listed scopes.",
+    )
+    audience: str | None = Field(
+        default=None,
+        description=(
+            "Expected audience (`aud`) claim for this API. If set, validation will reject tokens without this audience."
+        ),
+    )
+
+    # JWT verification params
+    jwks_uri: str | None = Field(
+        default=None,
+        description=("Direct JWKS endpoint URI for JWT signature verification. "
+                     "Optional if discovery or issuer is provided."),
+    )
+    discovery_url: str | None = Field(
+        default=None,
+        description=("OIDC discovery metadata URL. Used to automatically resolve JWKS and introspection endpoints."),
+    )
+
+    # Opaque token (introspection) params
+    introspection_endpoint: str | None = Field(
+        default=None,
+        description=("RFC 7662 token introspection endpoint. "
+                     "Required for opaque token validation and must be used with `client_id` and `client_secret`."),
+    )
+    client_id: str | None = Field(
+        default=None,
+        description="OAuth2 client ID for authenticating to the introspection endpoint (opaque token validation).",
+    )
+    client_secret: str | None = Field(
+        default=None,
+        description="OAuth2 client secret for authenticating to the introspection endpoint (opaque token validation).",
+    )
+
+    @staticmethod
+    def _is_https_or_localhost(url: str) -> bool:
+        try:
+            value = urlparse(url)
+            if not value.scheme or not value.netloc:
+                return False
+            if value.scheme == "https":
+                return True
+            return value.scheme == "http" and (value.hostname in {"localhost", "127.0.0.1", "::1"})
+        except Exception:
+            return False
+
+    @field_validator("issuer_url", "jwks_uri", "discovery_url", "introspection_endpoint")
+    @classmethod
+    def _require_valid_url(cls, value: str | None, info):
+        if value is None:
+            return value
+        if not cls._is_https_or_localhost(value):
+            raise ValueError(f"{info.field_name} must be HTTPS (http allowed only for localhost). Got: {value}")
+        return value
+
+    # ---------- Cross-field validation: ensure at least one viable path ----------
+
+    @model_validator(mode="after")
+    def _ensure_verification_path(self):
+        """
+        JWT path viable if any of: jwks_uri OR discovery_url OR issuer_url (fallback JWKS).
+        Opaque path viable if: introspection_endpoint AND client_id AND client_secret.
+        """
+        has_jwt_path = bool(self.jwks_uri or self.discovery_url or self.issuer_url)
+        has_opaque_path = bool(self.introspection_endpoint and self.client_id and self.client_secret)
+
+        # If introspection endpoint is set, enforce creds are present
+        if self.introspection_endpoint:
+            missing = []
+            if not self.client_id:
+                missing.append("client_id")
+            if not self.client_secret:
+                missing.append("client_secret")
+            if missing:
+                raise ValueError(
+                    f"introspection_endpoint configured but missing required credentials: {', '.join(missing)}")
+
+        # Require at least one path
+        if not (has_jwt_path or has_opaque_path):
+            raise ValueError("Invalid configuration: no verification method available. "
+                             "Configure one of the following:\n"
+                             "  • JWT path: set jwks_uri OR discovery_url OR issuer_url (for JWKS fallback)\n"
+                             "  • Opaque path: set introspection_endpoint + client_id + client_secret")
+
+        return self

nat/authentication/oauth2/register.py

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nat.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
+from nat.builder.builder import Builder
+from nat.cli.register_workflow import register_auth_provider
+
+
+@register_auth_provider(config_type=OAuth2AuthCodeFlowProviderConfig)
+async def oauth2_client(authentication_provider: OAuth2AuthCodeFlowProviderConfig, builder: Builder):
+    from nat.authentication.oauth2.oauth2_auth_code_flow_provider import OAuth2AuthCodeFlowProvider
+
+    yield OAuth2AuthCodeFlowProvider(authentication_provider)

nat/authentication/register.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# flake8: noqa
+
+from nat.authentication.api_key import register as register_api_key
+from nat.authentication.http_basic_auth import register as register_http_basic_auth
+from nat.authentication.oauth2 import register as register_oauth2