zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/zen_server/auth.py

@@ -14,13 +14,13 @@
 """Authentication module for ZenML server."""
 
 from contextvars import ContextVar
-from datetime import datetime
+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, HTTPException, status
+from fastapi import Depends, Response
 from fastapi.security import (
     HTTPBasic,
     HTTPBasicCredentials,
@@ -37,8 +37,17 @@ from zenml.constants import (
     LOGIN,
     VERSION_1,
 )
-from zenml.enums import AuthScheme, OAuthDeviceStatus
-from zenml.exceptions import AuthorizationException, OAuthError
+from zenml.enums import (
+    AuthScheme,
+    ExecutionStatus,
+    OAuthDeviceStatus,
+    OnboardingStep,
+)
+from zenml.exceptions import (
+    AuthorizationException,
+    CredentialsNotValid,
+    OAuthError,
+)
 from zenml.logger import get_logger
 from zenml.models import (
     APIKey,
@@ -47,13 +56,22 @@ from zenml.models import (
     ExternalUserModel,
     OAuthDeviceInternalResponse,
     OAuthDeviceInternalUpdate,
+    OAuthTokenResponse,
     UserAuthModel,
     UserRequest,
     UserResponse,
     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__)
 
@@ -109,7 +127,7 @@ def _fetch_and_verify_api_key(
         The fetched API key.
 
     Raises:
-        AuthorizationException: If the API key could not be found, is not
+        CredentialsNotValid: If the API key could not be found, is not
             active, if it could not be verified against the supplied key value
             or if the associated service account is not active.
     """
@@ -118,11 +136,9 @@ def _fetch_and_verify_api_key(
     try:
         api_key = zen_store().get_internal_api_key(api_key_id)
     except KeyError:
-        error = (
-            f"Authentication error: error retrieving API key " f"{api_key_id}"
-        )
+        error = f"Authentication error: error retrieving API key {api_key_id}"
         logger.error(error)
-        raise AuthorizationException(error)
+        raise CredentialsNotValid(error)
 
     if not api_key.service_account.active:
         error = (
@@ -131,7 +147,7 @@ def _fetch_and_verify_api_key(
             f"associated with API key {api_key.name} is not active"
         )
         logger.exception(error)
-        raise AuthorizationException(error)
+        raise CredentialsNotValid(error)
 
     if not api_key.active:
         error = (
@@ -141,7 +157,7 @@ def _fetch_and_verify_api_key(
             f"{api_key.service_account.name} is not active"
         )
         logger.error(error)
-        raise AuthorizationException(error)
+        raise CredentialsNotValid(error)
 
     if key_to_verify and not api_key.verify_key(key_to_verify):
         error = (
@@ -149,12 +165,12 @@ def _fetch_and_verify_api_key(
             f"{api_key.name}"
         )
         logger.exception(error)
-        raise AuthorizationException(error)
+        raise CredentialsNotValid(error)
 
     # Update the "last used" timestamp of the API key
     store.update_internal_api_key(
         api_key.id,
-        APIKeyInternalUpdate(update_last_login=True),  # type: ignore[call-arg]
+        APIKeyInternalUpdate(update_last_login=True),
     )
 
     return api_key
@@ -164,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.
@@ -182,13 +199,14 @@ 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:
         The authenticated account details.
 
     Raises:
-        AuthorizationException: If the credentials are invalid.
+        CredentialsNotValid: If the credentials are invalid.
     """
     user: Optional[UserAuthModel] = None
     auth_context: Optional[AuthContext] = None
@@ -218,11 +236,11 @@ def authenticate_credentials(
         if not UserAuthModel.verify_password(password, user):
             error = "Authentication error: invalid username or password"
             logger.error(error)
-            raise AuthorizationException(error)
+            raise CredentialsNotValid(error)
         if user and not user.active:
             error = f"Authentication error: user {user.name} is not active"
             logger.error(error)
-            raise AuthorizationException(error)
+            raise CredentialsNotValid(error)
 
     elif activation_token is not None:
         if not UserAuthModel.verify_activation_token(activation_token, user):
@@ -231,17 +249,33 @@ def authenticate_credentials(
                 f"{user_name_or_id}"
             )
             logger.error(error)
-            raise AuthorizationException(error)
+            raise CredentialsNotValid(error)
 
     elif access_token is not None:
         try:
             decoded_token = JWTToken.decode_token(
                 token=access_token,
             )
-        except AuthorizationException:
-            error = "Authentication error: error decoding access token. You may need to rerun zenml connect."
+        except CredentialsNotValid as e:
+            error = f"Authentication error: error decoding access token: {e}."
             logger.exception(error)
-            raise AuthorizationException(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(
@@ -253,7 +287,7 @@ def authenticate_credentials(
                 f"{decoded_token.user_id}"
             )
             logger.error(error)
-            raise AuthorizationException(error)
+            raise CredentialsNotValid(error)
 
         if not user_model.active:
             error = (
@@ -261,7 +295,7 @@ def authenticate_credentials(
                 f"active"
             )
             logger.error(error)
-            raise AuthorizationException(error)
+            raise CredentialsNotValid(error)
 
         api_key_model: Optional[APIKeyInternalResponse] = None
         if decoded_token.api_key_id:
@@ -272,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.
@@ -285,7 +327,7 @@ def authenticate_credentials(
                     f"{decoded_token.device_id}"
                 )
                 logger.error(error)
-                raise AuthorizationException(error)
+                raise CredentialsNotValid(error)
 
             if (
                 device_model.user is None
@@ -296,7 +338,7 @@ def authenticate_credentials(
                     f"does not belong to user {user_model.name}"
                 )
                 logger.error(error)
-                raise AuthorizationException(error)
+                raise CredentialsNotValid(error)
 
             if device_model.status != OAuthDeviceStatus.ACTIVE:
                 error = (
@@ -304,7 +346,7 @@ def authenticate_credentials(
                     f"is not active"
                 )
                 logger.error(error)
-                raise AuthorizationException(error)
+                raise CredentialsNotValid(error)
 
             if (
                 device_model.expires
@@ -315,7 +357,7 @@ def authenticate_credentials(
                     "has expired"
                 )
                 logger.error(error)
-                raise AuthorizationException(error)
+                raise CredentialsNotValid(error)
 
             zen_store().update_internal_authorized_device(
                 device_id=device_model.id,
@@ -324,6 +366,142 @@ def authenticate_credentials(
                 ),
             )
 
+        if decoded_token.schedule_id:
+            # If the token contains a schedule ID, we need to check if the
+            # schedule still exists in the database. We use a cached version
+            # of the schedule active status to avoid unnecessary database
+            # queries.
+
+            @cache_result(expiry=30)
+            def get_schedule_active(schedule_id: UUID) -> Optional[bool]:
+                """Get the active status of a schedule.
+
+                Args:
+                    schedule_id: The schedule ID.
+
+                Returns:
+                    The schedule active status or None if the schedule does not
+                    exist.
+                """
+                try:
+                    schedule = zen_store().get_schedule(
+                        schedule_id, hydrate=False
+                    )
+                except KeyError:
+                    return False
+
+                return schedule.active
+
+            schedule_active = get_schedule_active(decoded_token.schedule_id)
+            if schedule_active is None:
+                error = (
+                    f"Authentication error: error retrieving token schedule "
+                    f"{decoded_token.schedule_id}"
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+            if not schedule_active:
+                error = (
+                    f"Authentication error: schedule {decoded_token.schedule_id} "
+                    "is not active"
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+        if decoded_token.pipeline_run_id:
+            # If the token contains a pipeline run ID, we need to check if the
+            # pipeline run exists in the database and the pipeline run has
+            # not concluded. We use a cached version of the pipeline run status
+            # to avoid unnecessary database queries.
+
+            @cache_result(expiry=30)
+            def get_pipeline_run_status(
+                pipeline_run_id: UUID,
+            ) -> Optional[ExecutionStatus]:
+                """Get the status of a pipeline run.
+
+                Args:
+                    pipeline_run_id: The pipeline run ID.
+
+                Returns:
+                    The pipeline run status or None if the pipeline run does not
+                    exist.
+                """
+                try:
+                    pipeline_run = zen_store().get_run(
+                        pipeline_run_id, hydrate=False
+                    )
+                except KeyError:
+                    return None
+
+                return pipeline_run.status
+
+            pipeline_run_status = get_pipeline_run_status(
+                decoded_token.pipeline_run_id
+            )
+            if pipeline_run_status is None:
+                error = (
+                    f"Authentication error: error retrieving token pipeline run "
+                    f"{decoded_token.pipeline_run_id}"
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+            if pipeline_run_status.is_finished:
+                error = (
+                    f"The execution of pipeline run "
+                    f"{decoded_token.pipeline_run_id} has already concluded and "
+                    "API tokens scoped to it are no longer valid."
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+        if decoded_token.step_run_id:
+            # If the token contains a step run ID, we need to check if the
+            # step run exists in the database and the step run has not concluded.
+            # We use a cached version of the step run status to avoid unnecessary
+            # database queries.
+
+            @cache_result(expiry=30)
+            def get_step_run_status(
+                step_run_id: UUID,
+            ) -> Optional[ExecutionStatus]:
+                """Get the status of a step run.
+
+                Args:
+                    step_run_id: The step run ID.
+
+                Returns:
+                    The step run status or None if the step run does not exist.
+                """
+                try:
+                    step_run = zen_store().get_run_step(
+                        step_run_id, hydrate=False
+                    )
+                except KeyError:
+                    return None
+
+                return step_run.status
+
+            step_run_status = get_step_run_status(decoded_token.step_run_id)
+            if step_run_status is None:
+                error = (
+                    f"Authentication error: error retrieving token step run "
+                    f"{decoded_token.step_run_id}"
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+            if step_run_status.is_finished:
+                error = (
+                    f"The execution of step run "
+                    f"{decoded_token.step_run_id} has already concluded and "
+                    "API tokens scoped to it are no longer valid."
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
         auth_context = AuthContext(
             user=user_model,
             access_token=decoded_token,
@@ -340,12 +518,12 @@ def authenticate_credentials(
         if server_config().auth_scheme != AuthScheme.NO_AUTH:
             error = "Authentication error: no credentials provided"
             logger.error(error)
-            raise AuthorizationException(error)
+            raise CredentialsNotValid(error)
 
     if not auth_context:
         error = "Authentication error: invalid credentials"
         logger.error(error)
-        raise AuthorizationException(error)
+        raise CredentialsNotValid(error)
 
     return auth_context
 
@@ -487,12 +665,15 @@ def authenticate_device(client_id: UUID, device_code: str) -> AuthContext:
     return AuthContext(user=device_model.user, device=device_model)
 
 
-def authenticate_external_user(external_access_token: str) -> AuthContext:
+def authenticate_external_user(
+    external_access_token: str, request: Request
+) -> AuthContext:
     """Implement external authentication.
 
     Args:
         external_access_token: The access token used to authenticate the user
             to the external authenticator.
+        request: The request object.
 
     Returns:
         The authentication context reflecting the authenticated user.
@@ -511,6 +692,7 @@ def authenticate_external_user(external_access_token: str) -> AuthContext:
     # 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:
@@ -522,8 +704,7 @@ def authenticate_external_user(external_access_token: str) -> AuthContext:
         )
     except Exception as e:
         logger.exception(
-            f"Error fetching user information from external authenticator: "
-            f"{e}"
+            f"Error fetching user information from external authenticator: {e}"
         )
         raise AuthorizationException(
             "Error fetching user information from external authenticator."
@@ -544,7 +725,7 @@ def authenticate_external_user(external_access_token: str) -> AuthContext:
 
         if isinstance(payload, dict):
             try:
-                external_user = ExternalUserModel.parse_obj(payload)
+                external_user = ExternalUserModel.model_validate(payload)
             except Exception as e:
                 logger.exception(
                     f"Error parsing user information from external "
@@ -588,6 +769,7 @@ def authenticate_external_user(external_access_token: str) -> AuthContext:
                 email_opted_in=True,
                 active=True,
                 email=external_user.email,
+                is_admin=external_user.is_admin,
             ),
         )
     except KeyError:
@@ -603,15 +785,30 @@ def authenticate_external_user(external_access_token: str) -> AuthContext:
                 email_opted_in=True,
                 active=True,
                 email=external_user.email,
+                is_admin=external_user.is_admin,
             )
         )
 
         with AnalyticsContext() as context:
             context.user_id = user.id
             context.identify(
-                traits={"email": user.email, "source": "external_auth"}
+                traits={
+                    "email": external_user.email,
+                    "source": "external_auth",
+                }
             )
-            context.alias(user_id=user.id, previous_id=external_user.id)
+            context.alias(user_id=external_user.id, previous_id=user.id)
+
+    # This is the best spot to update the onboarding state to mark the
+    # "zenml login" step as completed for ZenML Pro servers, because the
+    # user has just successfully logged in. However, we need to differentiate
+    # between web clients (i.e. the dashboard) and CLI clients (i.e. the
+    # zenml CLI).
+    user_agent = request.headers.get("User-Agent", "").lower()
+    if "zenml/" in user_agent:
+        store.update_onboarding_state(
+            completed_steps={OnboardingStep.DEVICE_VERIFIED}
+        )
 
     return AuthContext(user=user)
 
@@ -629,14 +826,14 @@ def authenticate_api_key(
         The authentication context reflecting the authenticated service account.
 
     Raises:
-        AuthorizationException: If the service account could not be authorized.
+        CredentialsNotValid: If the service account could not be authorized.
     """
     try:
         decoded_api_key = APIKey.decode_api_key(api_key)
     except ValueError:
         error = "Authentication error: error decoding API key"
         logger.exception(error)
-        raise AuthorizationException(error)
+        raise CredentialsNotValid(error)
 
     internal_api_key = _fetch_and_verify_api_key(
         api_key_id=decoded_api_key.id, key_to_verify=decoded_api_key.key
@@ -650,6 +847,136 @@ def authenticate_api_key(
     return AuthContext(user=user_model, api_key=internal_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,
+    schedule_id: Optional[UUID] = None,
+    pipeline_run_id: Optional[UUID] = None,
+    step_run_id: Optional[UUID] = None,
+) -> OAuthTokenResponse:
+    """Generates an access token for the given user.
+
+    Args:
+        user_id: The ID of the user.
+        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,
+            the default value is determined automatically based on the server
+            configuration and type of token. If set to 0, the token will not
+            expire.
+        schedule_id: The ID of the schedule to scope the token to.
+        pipeline_run_id: The ID of the pipeline run to scope the token to.
+        step_run_id: The ID of the step run to scope the token to.
+
+    Returns:
+        An authentication response with an access token.
+    """
+    config = server_config()
+
+    # If the expiration time is not supplied, the JWT tokens are set to expire
+    # according to the values configured in the server config. Device tokens are
+    # handled separately from regular user tokens.
+    expires: Optional[datetime] = None
+    if expires_in == 0:
+        expires_in = None
+    elif expires_in is not None:
+        expires = datetime.utcnow() + timedelta(seconds=expires_in)
+    elif device:
+        # If a device was used for authentication, the token will expire
+        # at the same time as the device.
+        expires = device.expires
+        if expires:
+            expires_in = max(
+                int(expires.timestamp() - datetime.utcnow().timestamp()), 0
+            )
+    elif config.jwt_token_expire_minutes:
+        expires = datetime.utcnow() + timedelta(
+            minutes=config.jwt_token_expire_minutes
+        )
+        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,
+        api_key_id=api_key.id if api_key else None,
+        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 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,
+            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,
+            domain=config.auth_cookie_domain,
+        )
+
+    return OAuthTokenResponse(
+        access_token=access_token,
+        expires_in=expires_in,
+        token_type="bearer",
+        csrf_token=csrf_token,
+    )
+
+
 def http_authentication(
     credentials: HTTPBasicCredentials = Depends(HTTPBasic()),
 ) -> AuthContext:
@@ -661,19 +988,19 @@ def http_authentication(
     Returns:
         The authentication context reflecting the authenticated user.
 
-    Raises:
-        HTTPException: If the credentials are invalid.
+    # noqa: DAR401
     """
     try:
         return authenticate_credentials(
             user_name_or_id=credentials.username, password=credentials.password
         )
-    except AuthorizationException as e:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=str(e),
-            headers={"WWW-Authenticate": "Basic"},
-        )
+    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
+        # client can distinguish between a 401 error due to invalid credentials
+        # and other 401 errors and handle them accordingly by throwing away the
+        # current access token and re-authenticating.
+        raise http_exception_from_error(e)
 
 
 class CookieOAuth2TokenBearer(OAuth2PasswordBearer):
@@ -702,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,
@@ -712,21 +1040,25 @@ 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.
 
-    Raises:
-        HTTPException: If the JWT token could not be authorized.
+    # noqa: DAR401
     """
+    csrf_token = request.headers.get("X-CSRF-Token")
     try:
-        auth_context = authenticate_credentials(access_token=token)
-    except AuthorizationException as e:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=str(e),
-            headers={"WWW-Authenticate": "Bearer"},
+        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
+        # client can distinguish between a 401 error due to invalid credentials
+        # and other 401 errors and handle them accordingly by throwing away the
+        # current access token and re-authenticating.
+        raise http_exception_from_error(e)
 
     return auth_context