content: #  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.
  """HyperAI Service Connector.
  
  The HyperAI Service Connector allows authenticating to HyperAI (hyperai.ai)
  GPU equipped instances.
  """
  
  import base64
  import io
  from typing import Any, List, Optional, Type
  
  import paramiko
  from pydantic import Field
  
  from zenml.exceptions import AuthorizationException
  from zenml.integrations.hyperai import (
      HYPERAI_CONNECTOR_TYPE,
      HYPERAI_RESOURCE_TYPE,
  )
  from zenml.logger import get_logger
  from zenml.models import (
      AuthenticationMethodModel,
      ResourceTypeModel,
      ServiceConnectorTypeModel,
  )
  from zenml.service_connectors.service_connector import (
      AuthenticationConfig,
      ServiceConnector,
  )
  from zenml.utils.enum_utils import StrEnum
  from zenml.utils.secret_utils import PlainSerializedSecretStr
  
  logger = get_logger(__name__)
  
  
  class HyperAICredentials(AuthenticationConfig):
      """HyperAI client authentication credentials."""
  
      base64_ssh_key: PlainSerializedSecretStr = Field(
          title="SSH key (base64)",
      )
      ssh_passphrase: Optional[PlainSerializedSecretStr] = Field(
          default=None,
          title="SSH key passphrase",
      )
  
  
  class HyperAIConfiguration(HyperAICredentials):
      """HyperAI client configuration."""
  
      hostnames: List[str] = Field(
          title="Hostnames of the supported HyperAI instances.",
      )
  
      username: str = Field(
          title="Username to use to connect to the HyperAI instance.",
      )
  
  
  class HyperAIAuthenticationMethods(StrEnum):
      """HyperAI Authentication methods."""
  
      RSA_KEY_OPTIONAL_PASSPHRASE = "rsa-key"
      DSA_KEY_OPTIONAL_PASSPHRASE = "dsa-key"
      ECDSA_KEY_OPTIONAL_PASSPHRASE = "ecdsa-key"
      ED25519_KEY_OPTIONAL_PASSPHRASE = "ed25519-key"
  
  
  HYPERAI_SERVICE_CONNECTOR_TYPE_SPEC = ServiceConnectorTypeModel(
      name="HyperAI Service Connector",
      connector_type=HYPERAI_CONNECTOR_TYPE,
      description="""
  The ZenML HyperAI Service Connector allows authenticating to HyperAI (hyperai.ai)
  GPU equipped instances.
  
  This connector provides an SSH connection to your HyperAI instance, which can be
  used to run ZenML pipelines.
  
  The instance must be configured to allow SSH connections from the ZenML server.
  Docker and Docker Compose must be installed on the HyperAI instance. If you want
  to use scheduled pipeline runs, also ensure that a working cron daemon is installed
  and running on the HyperAI instance.
  """,
      logo_url="https://public-flavor-logos.s3.eu-central-1.amazonaws.com/connectors/hyperai/hyperai.png",
      emoji=":robot_face:",
      auth_methods=[
          AuthenticationMethodModel(
              name="RSA key with optional passphrase",
              auth_method=HyperAIAuthenticationMethods.RSA_KEY_OPTIONAL_PASSPHRASE,
              description="""
  Use an RSA private key to authenticate with a HyperAI instance. The key may be
  encrypted with a passphrase. If the key is encrypted, the passphrase must be
  provided. Make sure to provide the key as a Base64 encoded string.
  """,
              config_class=HyperAIConfiguration,
          ),
          AuthenticationMethodModel(
              name="DSA/DSS key with optional passphrase",
              auth_method=HyperAIAuthenticationMethods.DSA_KEY_OPTIONAL_PASSPHRASE,
              description="""
  Use a DSA/DSS private key to authenticate with a HyperAI instance. The key may be
  encrypted with a passphrase. If the key is encrypted, the passphrase must be
  provided. Make sure to provide the key as a Base64 encoded string.
  """,
              config_class=HyperAIConfiguration,
          ),
          AuthenticationMethodModel(
              name="ECDSA key with optional passphrase",
              auth_method=HyperAIAuthenticationMethods.ECDSA_KEY_OPTIONAL_PASSPHRASE,
              description="""
  Use an ECDSA private key to authenticate with a HyperAI instance. The key may be
  encrypted with a passphrase. If the key is encrypted, the passphrase must be
  provided. Make sure to provide the key as a Base64 encoded string.
  """,
              config_class=HyperAIConfiguration,
          ),
          AuthenticationMethodModel(
              name="Ed25519 key with optional passphrase",
              auth_method=HyperAIAuthenticationMethods.ED25519_KEY_OPTIONAL_PASSPHRASE,
              description="""
  Use an Ed25519 private key to authenticate with a HyperAI instance. The key may be
  encrypted with a passphrase. If the key is encrypted, the passphrase must be
  provided. Make sure to provide the key as a Base64 encoded string.
  """,
              config_class=HyperAIConfiguration,
          ),
      ],
      resource_types=[
          ResourceTypeModel(
              name="HyperAI instance",
              resource_type=HYPERAI_RESOURCE_TYPE,
              description="""
  Allows users to access a HyperAI instance as a resource. When used by
  connector consumers, they are provided a pre-authenticated SSH client
  instance.
  """,
              auth_methods=HyperAIAuthenticationMethods.values(),
              supports_instances=True,
              logo_url="https://public-flavor-logos.s3.eu-central-1.amazonaws.com/connectors/hyperai/hyperai.png",
              emoji=":robot_face:",
          ),
      ],
  )
  
  
  class HyperAIServiceConnector(ServiceConnector):
      """HyperAI service connector."""
  
      config: HyperAIConfiguration
  
      @classmethod
      def _get_connector_type(cls) -> ServiceConnectorTypeModel:
          """Get the service connector specification.
  
          Returns:
              The service connector specification.
          """
          return HYPERAI_SERVICE_CONNECTOR_TYPE_SPEC
  
      def _paramiko_key_type_given_auth_method(self) -> Type[paramiko.PKey]:
          """Get the Paramiko key type given the authentication method.
  
          Returns:
              The Paramiko key type.
  
          Raises:
              ValueError: If the authentication method is invalid.
          """
          mapping = {
              HyperAIAuthenticationMethods.RSA_KEY_OPTIONAL_PASSPHRASE: paramiko.RSAKey,
              HyperAIAuthenticationMethods.DSA_KEY_OPTIONAL_PASSPHRASE: paramiko.DSSKey,
              HyperAIAuthenticationMethods.ECDSA_KEY_OPTIONAL_PASSPHRASE: paramiko.ECDSAKey,
              HyperAIAuthenticationMethods.ED25519_KEY_OPTIONAL_PASSPHRASE: paramiko.Ed25519Key,
          }
  
          try:
              return mapping[HyperAIAuthenticationMethods(self.auth_method)]
          except KeyError:
              raise ValueError(
                  f"Invalid authentication method: {self.auth_method}"
              )
  
      def _create_paramiko_client(
          self, hostname: str
      ) -> paramiko.client.SSHClient:
          """Create a Paramiko SSH client based on the configuration.
  
          Args:
              hostname: The hostname of the HyperAI instance.
  
          Returns:
              A Paramiko SSH client.
  
          Raises:
              AuthorizationException: If the client cannot be created.
          """
          if self.config.ssh_passphrase is None:
              ssh_passphrase = None
          else:
              ssh_passphrase = self.config.ssh_passphrase.get_secret_value()
  
          # Connect to the HyperAI instance
          try:
              # Convert the SSH key from base64 to string
              base64_key_value = self.config.base64_ssh_key.get_secret_value()
              ssh_key = base64.b64decode(base64_key_value).decode("utf-8")
              paramiko_key = None
  
              with io.StringIO(ssh_key) as f:
                  paramiko_key = self._paramiko_key_type_given_auth_method().from_private_key(
                      f, password=ssh_passphrase
                  )
  
              # Trim whitespace from the IP address
              hostname = hostname.strip()
  
              paramiko_client = paramiko.client.SSHClient()
              paramiko_client.set_missing_host_key_policy(
                  paramiko.AutoAddPolicy()  # nosec
              )
              paramiko_client.connect(
                  hostname=hostname,
                  username=self.config.username,
                  pkey=paramiko_key,
                  timeout=30,
              )
  
              return paramiko_client
  
          except paramiko.ssh_exception.BadHostKeyException as e:
              logger.error("Bad host key: %s", e)
          except paramiko.ssh_exception.AuthenticationException as e:
              logger.error("Authentication failed: %s", e)
          except paramiko.ssh_exception.SSHException as e:
              logger.error(
                  "SSH error: %s. A common cause for this error is selection of the wrong key type in your service connector.",
                  e,
              )
          except Exception as e:
              logger.error(
                  "Unknown error while connecting to HyperAI instance: %s. Please check your network connection, IP address, and authentication details.",
                  e,
              )
  
          raise AuthorizationException(
              "Could not create SSH client for HyperAI instance."
          )
  
      def _authorize_client(self, hostname: str) -> None:
          """Verify that the client can authenticate with the HyperAI instance.
  
          Args:
              hostname: The hostname of the HyperAI instance.
          """
          logger.info("Verifying connection to HyperAI instance...")
  
          paramiko_client = self._create_paramiko_client(hostname)
          paramiko_client.close()
  
      def _connect_to_resource(
          self,
          **kwargs: Any,
      ) -> Any:
          """Connect to a HyperAI instance. Returns an authenticated SSH client.
  
          Args:
              kwargs: Additional implementation specific keyword arguments to pass
                  to the session or client constructor.
  
          Returns:
              An authenticated Paramiko SSH client.
          """
          logger.info("Connecting to HyperAI instance...")
          assert self.resource_id is not None
  
          paramiko_client = self._create_paramiko_client(self.resource_id)
          return paramiko_client
  
      def _configure_local_client(
          self,
          **kwargs: Any,
      ) -> None:
          """There is no local client for the HyperAI connector, so it does nothing.
  
          Args:
              kwargs: Additional implementation specific keyword arguments to pass
                  to the session or client constructor.
  
          Raises:
              NotImplementedError: If there is no local client for the HyperAI
                  connector.
          """
          raise NotImplementedError(
              "There is no local client for the HyperAI connector."
          )
  
      @classmethod
      def _auto_configure(
          cls,
          auth_method: Optional[str] = None,
          resource_type: Optional[str] = None,
          resource_id: Optional[str] = None,
          **kwargs: Any,
      ) -> "HyperAIServiceConnector":
          """Auto-configure the connector.
  
          Not supported by the HyperAI connector.
  
          Args:
              auth_method: The particular authentication method to use. If not
                  specified, the connector implementation must decide which
                  authentication method to use or raise an exception.
              resource_type: The type of resource to configure.
              resource_id: The ID of the resource to configure. The
                  implementation may choose to either require or ignore this
                  parameter if it does not support or detect an resource type that
                  supports multiple instances.
              kwargs: Additional implementation specific keyword arguments to use.
  
          Raises:
              NotImplementedError: If the connector auto-configuration fails or
                  is not supported.
          """
          raise NotImplementedError(
              "Auto-configuration is not supported by the HyperAI connector."
          )
  
      def _verify(
          self,
          resource_type: Optional[str] = None,
          resource_id: Optional[str] = None,
      ) -> List[str]:
          """Verify that a connection can be established to the HyperAI instance.
  
          Args:
              resource_type: The type of resource to verify. Must be set to the
                  Docker resource type.
              resource_id: The HyperAI instance to verify.
  
          Returns:
              The resource ID if the connection can be established.
  
          Raises:
              ValueError: If the resource ID is not in the list of configured
                  hostnames.
          """
          if resource_id:
              if resource_id not in self.config.hostnames:
                  raise ValueError(
                      f"The supplied hostname '{resource_id}' is not in the list "
                      f"of configured hostnames: {self.config.hostnames}. Please "
                      f"check your configuration."
                  )
              hostnames = [resource_id]
          else:
              hostnames = self.config.hostnames
  
          resources = []
          for hostname in hostnames:
              self._authorize_client(hostname)
              resources.append(hostname)
  
          return resources
  