content: #  Copyright (c) ZenML GmbH 2022. 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.
  """JWT utilities module for ZenML server."""
  
  from datetime import datetime, timedelta
  from typing import (
      Any,
      Dict,
      Optional,
      cast,
  )
  from uuid import UUID
  
  import jwt
  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 JWTToken(BaseModel):
      """Pydantic object representing a JWT token.
  
      Attributes:
          user_id: The id of the authenticated User.
          device_id: The id of the authenticated device.
          api_key_id: The id of the authenticated API key for which this token
              was issued.
          schedule_id: The id of the schedule for which the token was issued.
          pipeline_run_id: The id of the pipeline run for which the token was
              issued.
          deployment_id: The id of the deployment for which the token was
              issued.
          step_run_id: The id of the step run for which the token was
              issued.
          session_id: The id of the authenticated session (used for CSRF).
          claims: The original token claims.
      """
  
      user_id: UUID
      device_id: Optional[UUID] = None
      api_key_id: Optional[UUID] = None
      schedule_id: Optional[UUID] = None
      pipeline_run_id: Optional[UUID] = None
      deployment_id: Optional[UUID] = None
      session_id: Optional[UUID] = None
      claims: Dict[str, Any] = {}
  
      @classmethod
      def decode_token(
          cls,
          token: str,
          verify: bool = True,
      ) -> "JWTToken":
          """Decodes a JWT access token.
  
          Decodes a JWT access token and returns a `JWTToken` object with the
          information retrieved from its subject claim.
  
          Args:
              token: The encoded JWT token.
              verify: Whether to verify the signature of the token.
  
          Returns:
              The decoded JWT access token.
  
          Raises:
              CredentialsNotValid: If the token is invalid.
          """
          config = server_config()
  
          try:
              claims_data = jwt.decode(
                  token,
                  config.jwt_secret_key,
                  algorithms=[config.jwt_token_algorithm],
                  audience=config.get_jwt_token_audience(),
                  issuer=config.get_jwt_token_issuer(),
                  verify=verify,
                  leeway=timedelta(seconds=config.jwt_token_leeway_seconds),
              )
              claims = cast(Dict[str, Any], claims_data)
          except jwt.PyJWTError as e:
              raise CredentialsNotValid(f"Invalid JWT token: {e}") from e
  
          subject: str = claims.pop("sub", "")
          if not subject:
              raise CredentialsNotValid(
                  "Invalid JWT token: the subject claim is missing"
              )
  
          try:
              user_id = UUID(subject)
          except ValueError:
              raise CredentialsNotValid(
                  "Invalid JWT token: the subject claim is not a valid UUID"
              )
  
          device_id: Optional[UUID] = None
          if "device_id" in claims:
              try:
                  device_id = UUID(claims.pop("device_id"))
              except ValueError:
                  raise CredentialsNotValid(
                      "Invalid JWT token: the device_id claim is not a valid "
                      "UUID"
                  )
  
          api_key_id: Optional[UUID] = None
          if "api_key_id" in claims:
              try:
                  api_key_id = UUID(claims.pop("api_key_id"))
              except ValueError:
                  raise CredentialsNotValid(
                      "Invalid JWT token: the api_key_id claim is not a valid "
                      "UUID"
                  )
  
          schedule_id: Optional[UUID] = None
          if "schedule_id" in claims:
              try:
                  schedule_id = UUID(claims.pop("schedule_id"))
              except ValueError:
                  raise CredentialsNotValid(
                      "Invalid JWT token: the schedule_id claim is not a valid "
                      "UUID"
                  )
  
          pipeline_run_id: Optional[UUID] = None
          if "pipeline_run_id" in claims:
              try:
                  pipeline_run_id = UUID(claims.pop("pipeline_run_id"))
              except ValueError:
                  raise CredentialsNotValid(
                      "Invalid JWT token: the pipeline_run_id claim is not a valid "
                      "UUID"
                  )
  
          deployment_id: Optional[UUID] = None
          if "deployment_id" in claims:
              try:
                  deployment_id = UUID(claims.pop("deployment_id"))
              except ValueError:
                  raise CredentialsNotValid(
                      "Invalid JWT token: the deployment_id claim is not a valid "
                      "UUID"
                  )
  
          session_id: Optional[UUID] = None
          if "session_id" in claims:
              try:
                  session_id = UUID(claims.pop("session_id"))
              except ValueError:
                  raise CredentialsNotValid(
                      "Invalid JWT token: the session_id claim is not a valid "
                      "UUID"
                  )
  
          return JWTToken(
              user_id=user_id,
              device_id=device_id,
              api_key_id=api_key_id,
              schedule_id=schedule_id,
              pipeline_run_id=pipeline_run_id,
              deployment_id=deployment_id,
              session_id=session_id,
              claims=claims,
          )
  
      def encode(self, expires: Optional[datetime] = None) -> str:
          """Creates a JWT access token.
  
          Encodes, signs and returns a JWT access token.
  
          Args:
              expires: Datetime after which the token will expire. If not
                  provided, the JWT token will not be set to expire.
  
          Returns:
              The generated access token.
          """
          config = server_config()
  
          claims: Dict[str, Any] = self.claims.copy()
  
          claims["sub"] = str(self.user_id)
          claims["iss"] = config.get_jwt_token_issuer()
          claims["aud"] = config.get_jwt_token_audience()
  
          if expires:
              claims["exp"] = expires
          else:
              claims.pop("exp", None)
  
          if self.device_id:
              claims["device_id"] = str(self.device_id)
          if self.api_key_id:
              claims["api_key_id"] = str(self.api_key_id)
          if self.schedule_id:
              claims["schedule_id"] = str(self.schedule_id)
          if self.pipeline_run_id:
              claims["pipeline_run_id"] = str(self.pipeline_run_id)
          if self.deployment_id:
              claims["deployment_id"] = str(self.deployment_id)
          if self.session_id:
              claims["session_id"] = str(self.session_id)
  
          return jwt.encode(
              claims,
              config.jwt_secret_key,
              algorithm=config.jwt_token_algorithm,
          )
  