zenml-nightly 0.72.0.dev20250115__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/zen_server/auth.py

@@ -16,8 +16,8 @@
 from contextvars import ContextVar
 from datetime import datetime, timedelta
 from typing import Callable, Optional, Union
-from urllib.parse import urlencode
-from uuid import UUID
+from urllib.parse import urlencode, urlparse
+from uuid import UUID, uuid4
 
 import requests
 from fastapi import Depends, Response
@@ -63,9 +63,15 @@ from zenml.models import (
     UserUpdate,
 )
 from zenml.zen_server.cache import cache_result
+from zenml.zen_server.csrf import CSRFToken
 from zenml.zen_server.exceptions import http_exception_from_error
 from zenml.zen_server.jwt import JWTToken
-from zenml.zen_server.utils import server_config, zen_store
+from zenml.zen_server.utils import (
+    get_zenml_headers,
+    is_same_or_subdomain,
+    server_config,
+    zen_store,
+)
 
 logger = get_logger(__name__)
 
@@ -174,6 +180,7 @@ def authenticate_credentials(
     user_name_or_id: Optional[Union[str, UUID]] = None,
     password: Optional[str] = None,
     access_token: Optional[str] = None,
+    csrf_token: Optional[str] = None,
     activation_token: Optional[str] = None,
 ) -> AuthContext:
     """Verify if user authentication credentials are valid.
@@ -192,6 +199,7 @@ def authenticate_credentials(
         user_name_or_id: The username or user ID.
         password: The password.
         access_token: The access token.
+        csrf_token: The CSRF token.
         activation_token: The activation token.
 
     Returns:
@@ -253,6 +261,22 @@ def authenticate_credentials(
             logger.exception(error)
             raise CredentialsNotValid(error)
 
+        if decoded_token.session_id:
+            if not csrf_token:
+                error = "Authentication error: missing CSRF token"
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+            decoded_csrf_token = CSRFToken.decode_token(csrf_token)
+
+            if decoded_csrf_token.session_id != decoded_token.session_id:
+                error = (
+                    "Authentication error: CSRF token does not match the "
+                    "access token"
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
         try:
             user_model = zen_store().get_user(
                 user_name_or_id=decoded_token.user_id, include_private=True
@@ -282,6 +306,14 @@ def authenticate_credentials(
 
         device_model: Optional[OAuthDeviceInternalResponse] = None
         if decoded_token.device_id:
+            if server_config().auth_scheme in [
+                AuthScheme.NO_AUTH,
+                AuthScheme.EXTERNAL,
+            ]:
+                error = "Authentication error: device authorization is not supported."
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
             # Access tokens that have been issued for a device are only valid
             # for that device, so we need to check if the device ID matches any
             # of the valid devices in the database.
@@ -660,6 +692,7 @@ def authenticate_external_user(
     # Get the user information from the external authenticator
     user_info_url = config.external_user_info_url
     headers = {"Authorization": "Bearer " + external_access_token}
+    headers.update(get_zenml_headers())
     query_params = dict(server_id=str(config.get_external_server_id()))
 
     try:
@@ -817,6 +850,7 @@ def authenticate_api_key(
 def generate_access_token(
     user_id: UUID,
     response: Optional[Response] = None,
+    request: Optional[Request] = None,
     device: Optional[OAuthDeviceInternalResponse] = None,
     api_key: Optional[APIKeyInternalResponse] = None,
     expires_in: Optional[int] = None,
@@ -828,7 +862,11 @@ def generate_access_token(
 
     Args:
         user_id: The ID of the user.
-        response: The FastAPI response object.
+        response: The FastAPI response object. If passed, the access
+            token will also be set as an HTTP only cookie in the response.
+        request: The FastAPI request object. Used to determine the request
+            origin and to decide whether to use cross-site security measures for
+            the access token cookie.
         device: The device used for authentication.
         api_key: The service account API key used for authentication.
         expires_in: The number of seconds until the token expires. If not set,
@@ -866,6 +904,46 @@ def generate_access_token(
         )
         expires_in = config.jwt_token_expire_minutes * 60
 
+    # Figure out if this is a same-site request or a cross-site request
+    same_site = True
+    if response and request:
+        # Extract the origin domain from the request; use the referer as a
+        # fallback
+        origin_domain: Optional[str] = None
+        origin = request.headers.get("origin", request.headers.get("referer"))
+        if origin:
+            # If the request origin is known, we use it to determine whether
+            # this is a cross-site request and enable additional security
+            # measures.
+            origin_domain = urlparse(origin).netloc
+
+        server_domain: Optional[str] = config.auth_cookie_domain
+        # If the server's cookie domain is not explicitly set in the
+        # server's configuration, we use other sources to determine it:
+        #
+        # 1. the server's root URL, if set in the server's configuration
+        # 2. the X-Forwarded-Host header, if set by the reverse proxy
+        # 3. the request URL, if all else fails
+        if not server_domain and config.server_url:
+            server_domain = urlparse(config.server_url).netloc
+        if not server_domain:
+            server_domain = request.headers.get(
+                "x-forwarded-host", request.url.netloc
+            )
+
+        # Same-site requests can come from the same domain or from a
+        # subdomain of the domain used to issue cookies.
+        if origin_domain and server_domain:
+            same_site = is_same_or_subdomain(origin_domain, server_domain)
+
+    csrf_token: Optional[str] = None
+    session_id: Optional[UUID] = None
+    if not same_site:
+        # If responding to a cross-site login request, we need to generate and
+        # sign a CSRF token associated with the authentication session.
+        session_id = uuid4()
+        csrf_token = CSRFToken(session_id=session_id).encode()
+
     access_token = JWTToken(
         user_id=user_id,
         device_id=device.id if device else None,
@@ -873,15 +951,18 @@ def generate_access_token(
         schedule_id=schedule_id,
         pipeline_run_id=pipeline_run_id,
         step_run_id=step_run_id,
+        # Set the session ID if this is a cross-site request
+        session_id=session_id,
     ).encode(expires=expires)
 
-    if not device and response:
+    if response:
         # Also set the access token as an HTTP only cookie in the response
         response.set_cookie(
             key=config.get_auth_cookie_name(),
             value=access_token,
             httponly=True,
-            samesite="lax",
+            secure=not same_site,
+            samesite="lax" if same_site else "none",
             max_age=config.jwt_token_expire_minutes * 60
             if config.jwt_token_expire_minutes
             else None,
@@ -889,7 +970,10 @@ def generate_access_token(
         )
 
     return OAuthTokenResponse(
-        access_token=access_token, expires_in=expires_in, token_type="bearer"
+        access_token=access_token,
+        expires_in=expires_in,
+        token_type="bearer",
+        csrf_token=csrf_token,
     )
 
 
@@ -945,6 +1029,7 @@ class CookieOAuth2TokenBearer(OAuth2PasswordBearer):
 
 
 def oauth2_authentication(
+    request: Request,
     token: str = Depends(
         CookieOAuth2TokenBearer(
             tokenUrl=server_config().root_url_path + API + VERSION_1 + LOGIN,
@@ -955,14 +1040,18 @@ def oauth2_authentication(
 
     Args:
         token: The JWT bearer token to be authenticated.
+        request: The FastAPI request object.
 
     Returns:
         The authentication context reflecting the authenticated user.
 
     # noqa: DAR401
     """
+    csrf_token = request.headers.get("X-CSRF-Token")
     try:
-        auth_context = authenticate_credentials(access_token=token)
+        auth_context = authenticate_credentials(
+            access_token=token, csrf_token=csrf_token
+        )
     except CredentialsNotValid as e:
         # We want to be very explicit here and return a CredentialsNotValid
         # exception encoded as a 401 Unauthorized error encoded, so that the

zenml/zen_server/cloud_utils.py

@@ -1,114 +1,92 @@
 """Utils concerning anything concerning the cloud control plane backend."""
 
-import os
 from datetime import datetime, timedelta, timezone
 from typing import Any, Dict, Optional
 
 import requests
-from pydantic import BaseModel, ConfigDict, field_validator
 from requests.adapters import HTTPAdapter, Retry
 
+from zenml.config.server_config import ServerProConfiguration
 from zenml.exceptions import SubscriptionUpgradeRequiredError
-from zenml.zen_server.utils import server_config
-
-ZENML_CLOUD_RBAC_ENV_PREFIX = "ZENML_CLOUD_"
+from zenml.zen_server.utils import get_zenml_headers, server_config
 
 _cloud_connection: Optional["ZenMLCloudConnection"] = None
 
 
-class ZenMLCloudConfiguration(BaseModel):
-    """ZenML Pro RBAC configuration."""
-
-    api_url: str
-    oauth2_client_id: str
-    oauth2_client_secret: str
-    oauth2_audience: str
-
-    @field_validator("api_url")
-    @classmethod
-    def _strip_trailing_slashes_url(cls, url: str) -> str:
-        """Strip any trailing slashes on the API URL.
-
-        Args:
-            url: The API URL.
-
-        Returns:
-            The API URL with potential trailing slashes removed.
-        """
-        return url.rstrip("/")
-
-    @classmethod
-    def from_environment(cls) -> "ZenMLCloudConfiguration":
-        """Get the RBAC configuration from environment variables.
-
-        Returns:
-            The RBAC configuration.
-        """
-        env_config: Dict[str, Any] = {}
-        for k, v in os.environ.items():
-            if v == "":
-                continue
-            if k.startswith(ZENML_CLOUD_RBAC_ENV_PREFIX):
-                env_config[k[len(ZENML_CLOUD_RBAC_ENV_PREFIX) :].lower()] = v
-
-        return ZenMLCloudConfiguration(**env_config)
-
-    model_config = ConfigDict(
-        # Allow extra attributes from configs of previous ZenML versions to
-        # permit downgrading
-        extra="allow"
-    )
-
-
 class ZenMLCloudConnection:
     """Class to use for communication between server and control plane."""
 
     def __init__(self) -> None:
         """Initialize the RBAC component."""
-        self._config = ZenMLCloudConfiguration.from_environment()
+        self._config = ServerProConfiguration.get_server_config()
         self._session: Optional[requests.Session] = None
         self._token: Optional[str] = None
         self._token_expires_at: Optional[datetime] = None
 
-    def get(
-        self, endpoint: str, params: Optional[Dict[str, Any]]
+    def request(
+        self,
+        method: str,
+        endpoint: str,
+        params: Optional[Dict[str, Any]] = None,
+        data: Optional[Dict[str, Any]] = None,
     ) -> requests.Response:
-        """Send a GET request using the active session.
+        """Send a request using the active session.
 
         Args:
+            method: The HTTP method to use.
             endpoint: The endpoint to send the request to. This will be appended
                 to the base URL.
             params: Parameters to include in the request.
+            data: Data to include in the request.
 
         Raises:
+            SubscriptionUpgradeRequiredError: If the current subscription tier
+                is insufficient for the attempted operation.
             RuntimeError: If the request failed.
-            SubscriptionUpgradeRequiredError: In case the current subscription
-                tier is insufficient for the attempted operation.
 
         Returns:
             The response.
         """
         url = self._config.api_url + endpoint
 
-        response = self.session.get(url=url, params=params, timeout=7)
+        response = self.session.request(
+            method=method, url=url, params=params, json=data, timeout=7
+        )
         if response.status_code == 401:
-            # If we get an Unauthorized error from the API serer, we refresh the
-            # auth token and try again
+            # Refresh the auth token and try again
             self._clear_session()
-            response = self.session.get(url=url, params=params, timeout=7)
+            response = self.session.request(
+                method=method, url=url, params=params, json=data, timeout=7
+            )
 
         try:
             response.raise_for_status()
-        except requests.HTTPError:
+        except requests.HTTPError as e:
             if response.status_code == 402:
                 raise SubscriptionUpgradeRequiredError(response.json())
             else:
                 raise RuntimeError(
-                    f"Failed with the following error {response} {response.text}"
+                    f"Failed while trying to contact the central zenml pro "
+                    f"service: {e}"
                 )
 
         return response
 
+    def get(
+        self, endpoint: str, params: Optional[Dict[str, Any]]
+    ) -> requests.Response:
+        """Send a GET request using the active session.
+
+        Args:
+            endpoint: The endpoint to send the request to. This will be appended
+                to the base URL.
+            params: Parameters to include in the request.
+
+        Returns:
+            The response.
+        """
+        return self.request(method="GET", endpoint=endpoint, params=params)
+
     def post(
         self,
         endpoint: str,
@@ -123,33 +101,33 @@ class ZenMLCloudConnection:
             params: Parameters to include in the request.
             data: Data to include in the request.
 
-        Raises:
-            RuntimeError: If the request failed.
-
         Returns:
             The response.
         """
-        url = self._config.api_url + endpoint
-
-        response = self.session.post(
-            url=url, params=params, json=data, timeout=7
+        return self.request(
+            method="POST", endpoint=endpoint, params=params, data=data
         )
-        if response.status_code == 401:
-            # Refresh the auth token and try again
-            self._clear_session()
-            response = self.session.post(
-                url=url, params=params, json=data, timeout=7
-            )
 
-        try:
-            response.raise_for_status()
-        except requests.HTTPError as e:
-            raise RuntimeError(
-                f"Failed while trying to contact the central zenml pro "
-                f"service: {e}"
-            )
+    def patch(
+        self,
+        endpoint: str,
+        params: Optional[Dict[str, Any]] = None,
+        data: Optional[Dict[str, Any]] = None,
+    ) -> requests.Response:
+        """Send a PATCH request using the active session.
 
-        return response
+        Args:
+            endpoint: The endpoint to send the request to. This will be appended
+                to the base URL.
+            params: Parameters to include in the request.
+            data: Data to include in the request.
+
+        Returns:
+            The response.
+        """
+        return self.request(
+            method="PATCH", endpoint=endpoint, params=params, data=data
+        )
 
     @property
     def session(self) -> requests.Session:
@@ -169,6 +147,8 @@ class ZenMLCloudConnection:
             self._session = requests.Session()
             token = self._fetch_auth_token()
             self._session.headers.update({"Authorization": "Bearer " + token})
+            # Add the ZenML specific headers
+            self._session.headers.update(get_zenml_headers())
 
             retries = Retry(
                 total=5, backoff_factor=0.1, status_forcelist=[502, 504]
@@ -194,7 +174,7 @@ class ZenMLCloudConnection:
         self._token_expires_at = None
 
     def _fetch_auth_token(self) -> str:
-        """Fetch an auth token for the Cloud API from auth0.
+        """Fetch an auth token from the Cloud API.
 
         Raises:
             RuntimeError: If the auth token can't be fetched.
@@ -210,11 +190,14 @@ class ZenMLCloudConnection:
         ):
             return self._token
 
-        # Get an auth token from auth0
+        # Get an auth token from the Cloud API
         login_url = f"{self._config.api_url}/auth/login"
         headers = {"content-type": "application/x-www-form-urlencoded"}
+        # Add zenml specific headers to the request
+        headers.update(get_zenml_headers())
         payload = {
-            "client_id": self._config.oauth2_client_id,
+            # The client ID is the external server ID
+            "client_id": str(server_config().get_external_server_id()),
             "client_secret": self._config.oauth2_client_secret,
             "audience": self._config.oauth2_audience,
             "grant_type": "client_credentials",
@@ -225,7 +208,9 @@ class ZenMLCloudConnection:
             )
             response.raise_for_status()
         except Exception as e:
-            raise RuntimeError(f"Error fetching auth token from auth0: {e}")
+            raise RuntimeError(
+                f"Error fetching auth token from the Cloud API: {e}"
+            )
 
         json_response = response.json()
         access_token = json_response.get("access_token", "")
@@ -237,7 +222,9 @@ class ZenMLCloudConnection:
             or not expires_in
             or not isinstance(expires_in, int)
         ):
-            raise RuntimeError("Could not fetch auth token from auth0.")
+            raise RuntimeError(
+                "Could not fetch auth token from the Cloud API."
+            )
 
         self._token = access_token
         self._token_expires_at = datetime.now(timezone.utc) + timedelta(
@@ -259,3 +246,8 @@ def cloud_connection() -> ZenMLCloudConnection:
         _cloud_connection = ZenMLCloudConnection()
 
     return _cloud_connection
+
+
+def send_pro_tenant_status_update() -> None:
+    """Send a tenant status update to the Cloud API."""
+    cloud_connection().patch("/tenant_status")

zenml/zen_server/csrf.py

@@ -0,0 +1,91 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  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:
+#
+#       https://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.
+"""CSRF token utilities module for ZenML server."""
+
+from uuid import UUID
+
+from pydantic import BaseModel
+
+from zenml.exceptions import CredentialsNotValid
+from zenml.logger import get_logger
+from zenml.zen_server.utils import server_config
+
+logger = get_logger(__name__)
+
+
+class CSRFToken(BaseModel):
+    """Pydantic object representing a CSRF token.
+
+    Attributes:
+        session_id: The id of the authenticated session.
+    """
+
+    session_id: UUID
+
+    @classmethod
+    def decode_token(
+        cls,
+        token: str,
+    ) -> "CSRFToken":
+        """Decodes a CSRF token.
+
+        Decodes a CSRF access token and returns a `CSRFToken` object with the
+        information retrieved from its contents.
+
+        Args:
+            token: The encoded CSRF token.
+
+        Returns:
+            The decoded CSRF token.
+
+        Raises:
+            CredentialsNotValid: If the token is invalid.
+        """
+        from itsdangerous import BadData, BadSignature, URLSafeSerializer
+
+        config = server_config()
+
+        serializer = URLSafeSerializer(config.jwt_secret_key)
+        try:
+            # Decode and verify the token
+            data = serializer.loads(token)
+        except BadSignature as e:
+            raise CredentialsNotValid(
+                "Invalid CSRF token: signature mismatch"
+            ) from e
+        except BadData as e:
+            raise CredentialsNotValid("Invalid CSRF token") from e
+
+        try:
+            return CSRFToken(session_id=UUID(data))
+        except ValueError as e:
+            raise CredentialsNotValid(
+                "Invalid CSRF token: the session ID is not a valid UUID"
+            ) from e
+
+    def encode(self) -> str:
+        """Creates a CSRF token.
+
+        Encodes, signs and returns a CSRF access token.
+
+        Returns:
+            The generated CSRF token.
+        """
+        from itsdangerous import URLSafeSerializer
+
+        config = server_config()
+
+        serializer = URLSafeSerializer(config.jwt_secret_key)
+        token = serializer.dumps(str(self.session_id))
+        return token

zenml/zen_server/deploy/helm/templates/NOTES.txt

@@ -1,3 +1,24 @@
+{{- if .Values.zenml.pro.enabled }}
+
+The ZenML Pro server API is now active and ready to use at the following URL:
+
+  {{ .Values.zenml.serverURL }}
+
+{{- if .Values.zenml.pro.enrollmentKey }}
+
+The following enrollment key has been used to enroll your server in the ZenML Pro control plane:
+
+  {{ .Values.zenml.pro.enrollmentKey }}
+
+{{- else }}
+
+An enrollment key has been auto-generated for your server. Please use the following command to fetch the enrollment key:
+
+  kubectl get secret {{ include "zenml.fullname" . }} -o jsonpath="{.data.ZENML_SERVER_PRO_OAUTH2_CLIENT_SECRET}" | base64 --decode
+
+{{- end }}
+
+{{- else }}
 {{- if .Values.zenml.ingress.enabled }}
 {{- if .Values.zenml.ingress.host }}
 
@@ -28,3 +49,4 @@ You can get the ZenML server URL by running these commands:
 {{- end }}
 
 {{- end }}
+{{- end }}