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

zenml/zen_server/routers/auth_endpoints.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for authentication (login)."""
 
-from datetime import datetime, timedelta
 from typing import Optional, Union
 from urllib.parse import urlencode
 from uuid import UUID
@@ -40,16 +39,16 @@ from zenml.constants import (
     VERSION_1,
 )
 from zenml.enums import (
+    APITokenType,
     AuthScheme,
     OAuthDeviceStatus,
     OAuthGrantTypes,
 )
+from zenml.exceptions import AuthorizationException
 from zenml.logger import get_logger
 from zenml.models import (
-    APIKeyInternalResponse,
     OAuthDeviceAuthorizationResponse,
     OAuthDeviceInternalRequest,
-    OAuthDeviceInternalResponse,
     OAuthDeviceInternalUpdate,
     OAuthDeviceUserAgentHeader,
     OAuthRedirectResponse,
@@ -62,9 +61,10 @@ from zenml.zen_server.auth import (
     authenticate_device,
     authenticate_external_user,
     authorize,
+    generate_access_token,
 )
 from zenml.zen_server.exceptions import error_response
-from zenml.zen_server.jwt import JWTToken
+from zenml.zen_server.rate_limit import rate_limit_requests
 from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import verify_permission
 from zenml.zen_server.utils import (
@@ -114,6 +114,8 @@ class OAuthLoginRequestForm:
         Raises:
             HTTPException: If the request is invalid.
         """
+        config = server_config()
+
         if not grant_type:
             # Detect the grant type from the form data
             if username is not None:
@@ -122,8 +124,21 @@ class OAuthLoginRequestForm:
                 self.grant_type = OAuthGrantTypes.ZENML_API_KEY
             elif device_code:
                 self.grant_type = OAuthGrantTypes.OAUTH_DEVICE_CODE
-            else:
+            elif config.auth_scheme == AuthScheme.EXTERNAL:
                 self.grant_type = OAuthGrantTypes.ZENML_EXTERNAL
+            elif config.auth_scheme in [
+                AuthScheme.OAUTH2_PASSWORD_BEARER,
+                AuthScheme.NO_AUTH,
+                AuthScheme.HTTP_BASIC,
+            ]:
+                # For no auth and basic HTTP auth schemes, we also allow the
+                # password grant type to be used for backwards compatibility
+                self.grant_type = OAuthGrantTypes.OAUTH_PASSWORD
+            else:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail="Invalid request: grant type is required.",
+                )
         else:
             if grant_type not in OAuthGrantTypes.values():
                 logger.info(
@@ -135,10 +150,15 @@ class OAuthLoginRequestForm:
                 )
             self.grant_type = OAuthGrantTypes(grant_type)
 
-        config = server_config()
-
         if self.grant_type == OAuthGrantTypes.OAUTH_PASSWORD:
-            if config.auth_scheme != AuthScheme.OAUTH2_PASSWORD_BEARER:
+            # For the no auth and basic HTTP auth schemes, we also allow the
+            # password grant type to be used for compatibility with other
+            # auth schemes
+            if config.auth_scheme not in [
+                AuthScheme.OAUTH2_PASSWORD_BEARER,
+                AuthScheme.NO_AUTH,
+                AuthScheme.HTTP_BASIC,
+            ]:
                 logger.info(
                     f"Request with unsupported grant type: {self.grant_type}"
                 )
@@ -189,72 +209,14 @@ class OAuthLoginRequestForm:
                 )
 
 
-def generate_access_token(
-    user_id: UUID,
-    response: Response,
-    device: Optional[OAuthDeviceInternalResponse] = None,
-    api_key: Optional[APIKeyInternalResponse] = None,
-) -> OAuthTokenResponse:
-    """Generates an access token for the given user.
-
-    Args:
-        user_id: The ID of the user.
-        response: The FastAPI response object.
-        device: The device used for authentication.
-        api_key: The service account API key used for authentication.
-
-    Returns:
-        An authentication response with an access token.
-    """
-    config = server_config()
-
-    # 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
-    expires_in: Optional[int] = None
-    if 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
-
-    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,
-    ).encode(expires=expires)
-
-    if not device:
-        # 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",
-            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"
-    )
-
-
 @router.post(
     LOGIN,
     response_model=Union[OAuthTokenResponse, OAuthRedirectResponse],
 )
+@rate_limit_requests(
+    day_limit=server_config().login_rate_limit_day,
+    minute_limit=server_config().login_rate_limit_minute,
+)
 @handle_exceptions
 def token(
     request: Request,
@@ -274,6 +236,7 @@ def token(
     Raises:
         ValueError: If the grant type is invalid.
     """
+    config = server_config()
     if auth_form_data.grant_type == OAuthGrantTypes.OAUTH_PASSWORD:
         auth_context = authenticate_credentials(
             user_name_or_id=auth_form_data.username,
@@ -291,8 +254,6 @@ def token(
         )
 
     elif auth_form_data.grant_type == OAuthGrantTypes.ZENML_EXTERNAL:
-        config = server_config()
-
         assert config.external_cookie_name is not None
         assert config.external_login_url is not None
 
@@ -326,7 +287,8 @@ def token(
             return OAuthRedirectResponse(authorization_url=authorization_url)
 
         auth_context = authenticate_external_user(
-            external_access_token=external_access_token
+            external_access_token=external_access_token,
+            request=request,
         )
 
     else:
@@ -393,8 +355,11 @@ def device_authorization(
     config = server_config()
     store = zen_store()
 
-    # Use this opportunity to delete expired devices
-    store.delete_expired_authorized_devices()
+    try:
+        # Use this opportunity to delete expired devices
+        store.delete_expired_authorized_devices()
+    except Exception:
+        logger.exception("Failed to delete expired devices")
 
     # Fetch additional details about the client from the user-agent header
     user_agent_header = request.headers.get("User-Agent")
@@ -406,8 +371,14 @@ def device_authorization(
     # Fetch the IP address of the client
     ip_address: str = ""
     city, region, country = "", "", ""
-    if request.client and request.client.host:
+    forwarded = request.headers.get("X-Forwarded-For")
+
+    if forwarded:
+        ip_address = forwarded.split(",")[0].strip()
+    elif request.client and request.client.host:
         ip_address = request.client.host
+
+    if ip_address:
         city, region, country = get_ip_location(ip_address)
 
     # Check if a device is already registered for the same client ID.
@@ -424,7 +395,7 @@ def device_authorization(
                 city=city,
                 region=region,
                 country=country,
-                **device_details.dict(exclude_none=True),
+                **device_details.model_dump(exclude_none=True),
             )
         )
     else:
@@ -443,14 +414,14 @@ def device_authorization(
                 city=city,
                 region=region,
                 country=country,
-                **device_details.dict(exclude_none=True),
+                **device_details.model_dump(exclude_none=True),
             ),
         )
 
-    if config.dashboard_url:
-        verification_uri = (
-            config.dashboard_url.lstrip("/") + DEVICES + DEVICE_VERIFY
-        )
+    dashboard_url = config.dashboard_url or config.server_url
+
+    if dashboard_url:
+        verification_uri = dashboard_url.lstrip("/") + DEVICES + DEVICE_VERIFY
     else:
         verification_uri = DEVICES + DEVICE_VERIFY
 
@@ -480,52 +451,176 @@ def device_authorization(
 )
 @handle_exceptions
 def api_token(
-    pipeline_id: Optional[UUID] = None,
+    token_type: APITokenType = APITokenType.GENERIC,
+    expires_in: Optional[int] = None,
     schedule_id: Optional[UUID] = None,
-    expires_minutes: Optional[int] = None,
+    pipeline_run_id: Optional[UUID] = None,
+    step_run_id: Optional[UUID] = None,
     auth_context: AuthContext = Security(authorize),
 ) -> str:
-    """Get a workload API token for the current user.
+    """Generate an API token for the current user.
+
+    Use this endpoint to generate an API token for the current user. Two types
+    of API tokens are supported:
+
+    * Generic API token: This token is short-lived and can be used for
+    generic automation tasks. The expiration can be set by the user, but the
+    server will impose a maximum expiration time.
+    * Workload API token: This token is scoped to a specific pipeline run, step
+    run or schedule and is used by pipeline workloads to authenticate with the
+    server. A pipeline run ID, step run ID or schedule ID must be provided and
+    the generated token will only be valid for the indicated pipeline run, step
+    run or schedule. No time limit is imposed on the validity of the token.
+    A workload API token can be used to authenticate and generate another
+    workload API token, but only for the same schedule, pipeline run ID or step
+    run ID, in that order.
 
     Args:
-        pipeline_id: The ID of the pipeline to get the API token for.
-        schedule_id: The ID of the schedule to get the API token for.
-        expires_minutes: The number of minutes for which the API token should
-            be valid. If not provided, the API token will be valid indefinitely.
+        token_type: The type of API token to generate.
+        expires_in: The expiration time of the generic API token in seconds.
+            If not set, the server will use the default expiration time for
+            generic API tokens. The server also imposes a maximum expiration
+            time.
+        schedule_id: The ID of the schedule to scope the workload API token to.
+        pipeline_run_id: The ID of the pipeline run to scope the workload API
+            token to.
+        step_run_id: The ID of the step run to scope the workload API token to.
         auth_context: The authentication context.
 
     Returns:
         The API token.
 
     Raises:
-        HTTPException: If the user is not authenticated.
+        AuthorizationException: If not authorized to generate the API token.
+        ValueError: If the request is invalid.
     """
     token = auth_context.access_token
     if not token or not auth_context.encoded_access_token:
         # Should not happen
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Not authenticated.",
-        )
+        raise AuthorizationException("Not authenticated.")
+
+    if token_type == APITokenType.GENERIC:
+        if schedule_id or pipeline_run_id or step_run_id:
+            raise ValueError(
+                "Generic API tokens cannot be scoped to a schedule, pipeline "
+                "run or step run."
+            )
+
+        config = server_config()
+
+        if not expires_in:
+            expires_in = config.generic_api_token_lifetime
+
+        if expires_in > config.generic_api_token_max_lifetime:
+            raise ValueError(
+                f"The maximum expiration time for generic API tokens allowed "
+                f"by this server is {config.generic_api_token_max_lifetime} "
+                "seconds."
+            )
+
+        return generate_access_token(
+            user_id=token.user_id,
+            expires_in=expires_in,
+        ).access_token
 
     verify_permission(
         resource_type=ResourceType.PIPELINE_RUN, action=Action.CREATE
     )
 
-    if not token.device_id:
-        # If not authenticated with a device, the current API token is returned
-        # as is, without any modifications. Issuing workload tokens is only
-        # supported for device authenticated users, because device tokens can
-        # be revoked at any time.
-        return auth_context.encoded_access_token
-
-    # If authenticated with a device, a new API token is generated for the
-    # pipeline and/or schedule.
-    if pipeline_id:
-        token.pipeline_id = pipeline_id
+    schedule_id = schedule_id or token.schedule_id
+    pipeline_run_id = pipeline_run_id or token.pipeline_run_id
+    step_run_id = step_run_id or token.step_run_id
+
+    if not pipeline_run_id and not schedule_id and not step_run_id:
+        raise ValueError(
+            "Workload API tokens must be scoped to a schedule, pipeline run "
+            "or step run."
+        )
+
+    if schedule_id and token.schedule_id and schedule_id != token.schedule_id:
+        raise AuthorizationException(
+            f"Unable to scope API token to schedule {schedule_id}. The "
+            f"token used to authorize this request is already scoped to "
+            f"schedule {token.schedule_id}."
+        )
+
+    if (
+        pipeline_run_id
+        and token.pipeline_run_id
+        and pipeline_run_id != token.pipeline_run_id
+    ):
+        raise AuthorizationException(
+            f"Unable to scope API token to pipeline run {pipeline_run_id}. The "
+            f"token used to authorize this request is already scoped to "
+            f"pipeline run {token.pipeline_run_id}."
+        )
+
+    if step_run_id and token.step_run_id and step_run_id != token.step_run_id:
+        raise AuthorizationException(
+            f"Unable to scope API token to step run {step_run_id}. The "
+            f"token used to authorize this request is already scoped to "
+            f"step run {token.step_run_id}."
+        )
+
     if schedule_id:
-        token.schedule_id = schedule_id
-    expires: Optional[datetime] = None
-    if expires_minutes:
-        expires = datetime.utcnow() + timedelta(minutes=expires_minutes)
-    return token.encode(expires=expires)
+        # The schedule must exist
+        try:
+            schedule = zen_store().get_schedule(schedule_id, hydrate=False)
+        except KeyError:
+            raise ValueError(
+                f"Schedule {schedule_id} does not exist and API tokens cannot "
+                "be generated for non-existent schedules for security reasons."
+            )
+
+        if not schedule.active:
+            raise ValueError(
+                f"Schedule {schedule_id} is not active and API tokens cannot "
+                "be generated for inactive schedules for security reasons."
+            )
+
+    if pipeline_run_id:
+        # The pipeline run must exist and the run must not be concluded
+        try:
+            pipeline_run = zen_store().get_run(pipeline_run_id, hydrate=False)
+        except KeyError:
+            raise ValueError(
+                f"Pipeline run {pipeline_run_id} does not exist and API tokens "
+                "cannot be generated for non-existent pipeline runs for "
+                "security reasons."
+            )
+
+        if pipeline_run.status.is_finished:
+            raise ValueError(
+                f"The execution of pipeline run {pipeline_run_id} has already "
+                "concluded and API tokens can no longer be generated for it "
+                "for security reasons."
+            )
+
+    if step_run_id:
+        # The step run must exist and the step must not be concluded
+        try:
+            step_run = zen_store().get_run_step(step_run_id, hydrate=False)
+        except KeyError:
+            raise ValueError(
+                f"Step run {step_run_id} does not exist and API tokens cannot "
+                "be generated for non-existent step runs for security reasons."
+            )
+
+        if step_run.status.is_finished:
+            raise ValueError(
+                f"The execution of step run {step_run_id} has already "
+                "concluded and API tokens can no longer be generated for it "
+                "for security reasons."
+            )
+
+    return generate_access_token(
+        user_id=token.user_id,
+        # Keep the original API key and device token scopes
+        api_key=auth_context.api_key,
+        device=auth_context.device,
+        schedule_id=schedule_id,
+        pipeline_run_id=pipeline_run_id,
+        step_run_id=step_run_id,
+        # Never expire the token
+        expires_in=0,
+    ).access_token

zenml/zen_server/routers/code_repositories_endpoints.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Endpoint definitions for code repositories."""
+
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security

zenml/zen_server/routers/devices_endpoints.py

@@ -12,19 +12,22 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Endpoint definitions for code repositories."""
+
 from datetime import datetime
 from typing import Optional
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security
 
+from zenml.analytics.enums import AnalyticsEvent
+from zenml.analytics.utils import track_handler
 from zenml.constants import (
     API,
     DEVICE_VERIFY,
     DEVICES,
     VERSION_1,
 )
-from zenml.enums import OAuthDeviceStatus
+from zenml.enums import OAuthDeviceStatus, OnboardingStep
 from zenml.models import (
     OAuthDeviceFilter,
     OAuthDeviceInternalUpdate,
@@ -208,64 +211,68 @@ def verify_authorized_device(
         device_id=device_id,
     )
 
-    # Check if the device is in a state that allows verification.
-    if device_model.status != OAuthDeviceStatus.PENDING:
-        raise ValueError(
-            "Invalid request: device not pending verification.",
-        )
+    with track_handler(event=AnalyticsEvent.DEVICE_VERIFIED):
+        # Check if the device is in a state that allows verification.
+        if device_model.status != OAuthDeviceStatus.PENDING:
+            raise ValueError(
+                "Invalid request: device not pending verification.",
+            )
 
-    # Check if the device verification has expired.
-    if device_model.expires and device_model.expires < datetime.utcnow():
-        raise ValueError(
-            "Invalid request: device verification expired.",
-        )
+        # Check if the device verification has expired.
+        if device_model.expires and device_model.expires < datetime.utcnow():
+            raise ValueError(
+                "Invalid request: device verification expired.",
+            )
 
-    # Check if the device already has a user associated with it. If so, the
-    # current user and the user associated with the device must be the same.
-    if device_model.user and device_model.user.id != auth_context.user.id:
-        raise ValueError(
-            "Invalid request: this device is associated with another user.",
-        )
+        # Check if the device already has a user associated with it. If so, the
+        # current user and the user associated with the device must be the same.
+        if device_model.user and device_model.user.id != auth_context.user.id:
+            raise ValueError(
+                "Invalid request: this device is associated with another user.",
+            )
 
-    # Check if the device code is valid.
-    if not device_model.verify_user_code(request.user_code):
-        # If the device code is invalid, increment the failed auth attempts
-        # counter and lock the device if the maximum number of failed auth
-        # attempts has been reached.
-        failed_auth_attempts = device_model.failed_auth_attempts + 1
+        # Check if the device code is valid.
+        if not device_model.verify_user_code(request.user_code):
+            # If the device code is invalid, increment the failed auth attempts
+            # counter and lock the device if the maximum number of failed auth
+            # attempts has been reached.
+            failed_auth_attempts = device_model.failed_auth_attempts + 1
+            update = OAuthDeviceInternalUpdate(
+                failed_auth_attempts=failed_auth_attempts
+            )
+            if failed_auth_attempts >= config.max_failed_device_auth_attempts:
+                update.locked = True
+            store.update_internal_authorized_device(
+                device_id=device_model.id,
+                update=update,
+            )
+            if failed_auth_attempts >= config.max_failed_device_auth_attempts:
+                raise ValueError(
+                    "Invalid request: device locked due to too many failed "
+                    "authentication attempts.",
+                )
+            raise ValueError(
+                "Invalid request: invalid user code.",
+            )
+
+        # If the device code is valid, associate the device with the current user.
+        # We don't reset the expiration date yet, because we want to make sure
+        # that the client has received the access token before we do so, to avoid
+        # brute force attacks on the device code.
         update = OAuthDeviceInternalUpdate(
-            failed_auth_attempts=failed_auth_attempts
+            status=OAuthDeviceStatus.VERIFIED,
+            user_id=auth_context.user.id,
+            failed_auth_attempts=0,
+            trusted_device=request.trusted_device,
         )
-        if failed_auth_attempts >= config.max_failed_device_auth_attempts:
-            update.locked = True
-        store.update_internal_authorized_device(
+        device_model = store.update_internal_authorized_device(
             device_id=device_model.id,
             update=update,
         )
-        if failed_auth_attempts >= config.max_failed_device_auth_attempts:
-            raise ValueError(
-                "Invalid request: device locked due to too many failed "
-                "authentication attempts.",
-            )
-        raise ValueError(
-            "Invalid request: invalid user code.",
-        )
 
-    # If the device code is valid, associate the device with the current user.
-    # We don't reset the expiration date yet, because we want to make sure
-    # that the client has received the access token before we do so, to avoid
-    # brute force attacks on the device code.
-    update = OAuthDeviceInternalUpdate(
-        status=OAuthDeviceStatus.VERIFIED,
-        user_id=auth_context.user.id,
-        failed_auth_attempts=0,
-        trusted_device=request.trusted_device,
+    store.update_onboarding_state(
+        completed_steps={OnboardingStep.DEVICE_VERIFIED}
     )
-    device_model = store.update_internal_authorized_device(
-        device_id=device_model.id,
-        update=update,
-    )
-
     return device_model