client.py

# Copyright 2024 Google LLC
#
# 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
#
#     http://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.

from __future__ import annotations

import asyncio
import datetime
import logging
from typing import Any, Optional, TYPE_CHECKING

import aiohttp
from cryptography.hazmat.backends import default_backend
from cryptography.x509 import load_pem_x509_certificate
from google.auth.credentials import TokenState
from google.auth.transport import requests

from google.cloud.sql.connector.connection_info import ConnectionInfo
from google.cloud.sql.connector.exceptions import AutoIAMAuthNotSupported
from google.cloud.sql.connector.refresh_utils import _downscope_credentials
from google.cloud.sql.connector.refresh_utils import retry_50x
from google.cloud.sql.connector.version import __version__ as version

if TYPE_CHECKING:
    from google.auth.credentials import Credentials

USER_AGENT: str = f"cloud-sql-python-connector/{version}"
API_VERSION: str = "v1beta4"
DEFAULT_SERVICE_ENDPOINT: str = "https://sqladmin.googleapis.com"

logger = logging.getLogger(name=__name__)


def _format_user_agent(driver: Optional[str], custom: Optional[str]) -> str:
    agent = f"{USER_AGENT}+{driver}" if driver else USER_AGENT
    if custom and isinstance(custom, str):
        agent = f"{agent} {custom}"
    return agent


class CloudSQLClient:
    def __init__(
        self,
        sqladmin_api_endpoint: Optional[str],
        quota_project: Optional[str],
        credentials: Credentials,
        client: Optional[aiohttp.ClientSession] = None,
        driver: Optional[str] = None,
        user_agent: Optional[str] = None,
    ) -> None:
        """
        Establish the client to be used for Cloud SQL Admin API requests.

        Args:
            sqladmin_api_endpoint (str): Base URL to use when calling
                the Cloud SQL Admin API endpoints.
            quota_project (str): The Project ID for an existing Google Cloud
                project. The project specified is used for quota and
                billing purposes.
            credentials (google.auth.credentials.Credentials):
                A credentials object created from the google-auth Python library.
                Must have the Cloud SQL Admin scopes. For more info check out
                https://google-auth.readthedocs.io/en/latest/.
            client (aiohttp.ClientSession): Async client used to make requests to
                Cloud SQL Admin APIs.
                Optional, defaults to None and creates new client.
            driver (str): Database driver to be used by the client.
        """
        user_agent = _format_user_agent(driver, user_agent)
        headers = {
            "x-goog-api-client": user_agent,
            "User-Agent": user_agent,
            "Content-Type": "application/json",
        }
        if quota_project:
            headers["x-goog-user-project"] = quota_project

        self._client = client if client else aiohttp.ClientSession(headers=headers)
        self._credentials = credentials
        if sqladmin_api_endpoint is None:
            self._sqladmin_api_endpoint = DEFAULT_SERVICE_ENDPOINT
        else:
            self._sqladmin_api_endpoint = sqladmin_api_endpoint
        self._user_agent = user_agent

    async def _get_metadata(
        self,
        project: str,
        region: str,
        instance: str,
    ) -> dict[str, Any]:
        """Requests metadata from the Cloud SQL Instance
        and returns a dictionary containing the IP addresses and certificate
        authority of the Cloud SQL Instance.

        :type project: str
        :param project:
            A string representing the name of the project.

        :type region: str
        :param region : A string representing the name of the region.

        :type instance: str
        :param instance: A string representing the name of the instance.

        :rtype: dict[str: Union[dict, str]]
        :returns: Returns a dictionary containing a dictionary of all IP
            addresses and their type and a string representing the
            certificate authority.
        """

        headers = {
            "Authorization": f"Bearer {self._credentials.token}",
        }

        url = f"{self._sqladmin_api_endpoint}/sql/{API_VERSION}/projects/{project}/instances/{instance}/connectSettings"

        resp = await self._client.get(url, headers=headers)
        if resp.status >= 500:
            resp = await retry_50x(self._client.get, url, headers=headers)
        resp.raise_for_status()
        ret_dict = await resp.json()

        if ret_dict["region"] != region:
            raise ValueError(
                f'[{project}:{region}:{instance}]: Provided region was mismatched - got region {region}, expected {ret_dict["region"]}.'
            )

        ip_addresses = (
            {ip["type"]: ip["ipAddress"] for ip in ret_dict["ipAddresses"]}
            if "ipAddresses" in ret_dict
            else {}
        )
        # resolve dnsName into IP address for PSC
        # Note that we have to check for PSC enablement also because CAS
        # instances also set the dnsName field.
        # Remove trailing period from DNS name. Required for SSL in Python
        dns_name = ret_dict.get("dnsName", "").rstrip(".")
        if dns_name and ret_dict.get("pscEnabled"):
            ip_addresses["PSC"] = dns_name

        return {
            "ip_addresses": ip_addresses,
            "server_ca_cert": ret_dict["serverCaCert"]["cert"],
            "database_version": ret_dict["databaseVersion"],
        }

    async def _get_ephemeral(
        self,
        project: str,
        instance: str,
        pub_key: str,
        enable_iam_auth: bool = False,
    ) -> tuple[str, datetime.datetime]:
        """Asynchronously requests an ephemeral certificate from the Cloud SQL Instance.

        :type project: str
        :param project : A string representing the name of the project.

        :type instance: str
        :param instance: A string representing the name of the instance.

        :type pub_key:
        :param str: A string representing PEM-encoded RSA public key.

        :type enable_iam_auth: bool
        :param enable_iam_auth
            Enables automatic IAM database authentication for Postgres or MySQL
            instances.

        :rtype: str
        :returns: An ephemeral certificate from the Cloud SQL instance that allows
            authorized connections to the instance.
        """
        headers = {
            "Authorization": f"Bearer {self._credentials.token}",
        }

        url = f"{self._sqladmin_api_endpoint}/sql/{API_VERSION}/projects/{project}/instances/{instance}:generateEphemeralCert"

        data = {"public_key": pub_key}

        if enable_iam_auth:
            # down-scope credentials with only IAM login scope (refreshes them too)
            login_creds = _downscope_credentials(self._credentials)
            data["access_token"] = login_creds.token

        resp = await self._client.post(url, headers=headers, json=data)
        if resp.status >= 500:
            resp = await retry_50x(self._client.post, url, headers=headers, json=data)
        resp.raise_for_status()
        ret_dict = await resp.json()

        ephemeral_cert: str = ret_dict["ephemeralCert"]["cert"]

        # decode cert to read expiration
        x509 = load_pem_x509_certificate(
            ephemeral_cert.encode("UTF-8"), default_backend()
        )
        expiration = x509.not_valid_after_utc
        # for IAM authentication OAuth2 token is embedded in cert so it
        # must still be valid for successful connection
        if enable_iam_auth:
            token_expiration: datetime.datetime = login_creds.expiry
            # google.auth library strips timezone info for backwards compatibality
            # reasons with Python 2. Add it back to allow timezone aware datetimes.
            # Ref: https://github.com/googleapis/google-auth-library-python/blob/49a5ff7411a2ae4d32a7d11700f9f961c55406a9/google/auth/_helpers.py#L93-L99
            token_expiration = token_expiration.replace(tzinfo=datetime.timezone.utc)

            if expiration > token_expiration:
                expiration = token_expiration
        return ephemeral_cert, expiration

    async def get_connection_info(
        self,
        project: str,
        region: str,
        instance: str,
        keys: asyncio.Future,
        enable_iam_auth: bool,
    ) -> ConnectionInfo:
        """Immediately performs a full refresh operation using the Cloud SQL
        Admin API.

        Args:
            project (str): The name of the project the Cloud SQL instance is
                located in.
            region (str): The region the Cloud SQL instance is located in.
            instance (str): Name of the Cloud SQL instance.
            keys (asyncio.Future): A future to the client's public-private key
                pair.
            enable_iam_auth (bool): Whether an automatic IAM database
                authentication connection is being requested (Postgres and MySQL).

        Returns:
            ConnectionInfo: All the information required to connect securely to
                the Cloud SQL instance.
        Raises:
            AutoIAMAuthNotSupported: Database engine does not support automatic
                IAM authentication.
        """
        priv_key, pub_key = await keys
        # before making Cloud SQL Admin API calls, refresh creds if required
        if not self._credentials.token_state == TokenState.FRESH:
            self._credentials.refresh(requests.Request())

        metadata_task = asyncio.create_task(
            self._get_metadata(
                project,
                region,
                instance,
            )
        )

        ephemeral_task = asyncio.create_task(
            self._get_ephemeral(
                project,
                instance,
                pub_key,
                enable_iam_auth,
            )
        )
        try:
            metadata = await metadata_task
            # check if automatic IAM database authn is supported for database engine
            if enable_iam_auth and not metadata["database_version"].startswith(
                ("POSTGRES", "MYSQL")
            ):
                raise AutoIAMAuthNotSupported(
                    f"'{metadata['database_version']}' does not support "
                    "automatic IAM authentication. It is only supported with "
                    "Cloud SQL Postgres or MySQL instances."
                )
        except Exception:
            # cancel ephemeral cert task if exception occurs before it is awaited
            ephemeral_task.cancel()
            raise

        ephemeral_cert, expiration = await ephemeral_task

        return ConnectionInfo(
            ephemeral_cert,
            metadata["server_ca_cert"],
            priv_key,
            metadata["ip_addresses"],
            metadata["database_version"],
            expiration,
        )

    async def close(self) -> None:
        """Close CloudSQLClient gracefully."""
        logger.debug("Waiting for Connector's http client to close")
        await self._client.close()
        logger.debug("Closed Connector's http client")