oauth.py

import base64
import functools
import hashlib
import json
import logging
import os
import secrets
import threading
import urllib.parse
import webbrowser
from abc import abstractmethod
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from datetime import datetime, timedelta
from enum import Enum
from http.server import BaseHTTPRequestHandler, HTTPServer
from typing import Any, Callable, Dict, List, Optional

import requests
import requests.auth

from ._base_client import _BaseClient, _fix_host_if_needed
from .environments import Cloud

# Error code for PKCE flow in Azure Active Directory, that gets additional retry.
# See https://stackoverflow.com/a/75466778/277035 for more info
NO_ORIGIN_FOR_SPA_CLIENT_ERROR = "AADSTS9002327"

URL_ENCODED_CONTENT_TYPE = "application/x-www-form-urlencoded"
JWT_BEARER_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:jwt-bearer"
OIDC_TOKEN_PATH = "/oidc/v1/token"

logger = logging.getLogger(__name__)


@dataclass
class AuthorizationDetail:
    type: str
    object_type: str
    object_path: str
    actions: List[str]

    def as_dict(self) -> dict:
        return {
            "type": self.type,
            "object_type": self.object_type,
            "object_path": self.object_path,
            "actions": self.actions,
        }

    def from_dict(self, d: dict) -> "AuthorizationDetail":
        return AuthorizationDetail(
            type=d.get("type"),
            object_type=d.get("object_type"),
            object_path=d.get("object_path"),
            actions=d.get("actions"),
        )


class IgnoreNetrcAuth(requests.auth.AuthBase):
    """This auth method is a no-op.

    We use it to force requestslib to not use .netrc to write auth headers
    when making .post() requests to the oauth token endpoints, since these
    don't require authentication.

    In cases where .netrc is outdated or corrupt, these requests will fail.

    See issue #121
    """

    def __call__(self, r):
        return r


@dataclass
class OidcEndpoints:
    """
    The endpoints used for OAuth-based authentication in Databricks.
    """

    authorization_endpoint: str  # ../v1/authorize
    """The authorization endpoint for the OAuth flow. The user-agent should be directed to this endpoint in order for
    the user to login and authorize the client for user-to-machine (U2M) flows."""

    token_endpoint: str  # ../v1/token
    """The token endpoint for the OAuth flow."""

    @staticmethod
    def from_dict(d: dict) -> "OidcEndpoints":
        return OidcEndpoints(
            authorization_endpoint=d.get("authorization_endpoint"),
            token_endpoint=d.get("token_endpoint"),
        )

    def as_dict(self) -> dict:
        return {
            "authorization_endpoint": self.authorization_endpoint,
            "token_endpoint": self.token_endpoint,
        }


@dataclass
class Token:
    access_token: str
    token_type: Optional[str] = None
    refresh_token: Optional[str] = None
    expiry: Optional[datetime] = None

    @property
    def expired(self):
        if not self.expiry:
            return False
        # Azure Databricks rejects tokens that expire in 30 seconds or less,
        # so we refresh the token 40 seconds before it expires.
        potentially_expired = self.expiry - timedelta(seconds=40)
        now = datetime.now(tz=potentially_expired.tzinfo)
        is_expired = potentially_expired < now
        return is_expired

    @property
    def valid(self):
        return self.access_token and not self.expired

    def as_dict(self) -> dict:
        raw = {
            "access_token": self.access_token,
            "token_type": self.token_type,
        }
        if self.expiry:
            raw["expiry"] = self.expiry.isoformat()
        if self.refresh_token:
            raw["refresh_token"] = self.refresh_token
        return raw

    @staticmethod
    def from_dict(raw: dict) -> "Token":
        return Token(
            access_token=raw["access_token"],
            token_type=raw["token_type"],
            expiry=datetime.fromisoformat(raw["expiry"]),
            refresh_token=raw.get("refresh_token"),
        )

    def jwt_claims(self) -> Dict[str, str]:
        """Get claims from the access token or return an empty dictionary if it is not a JWT token.

        All refreshable tokens we're dealing with are JSON Web Tokens (JWT).

        The common claims are:
        - 'aud' represents the intended recipient of the token. In case of Azure, this is an app's Application ID
                assigned within the Azure portal.
        - 'iss' serves to identify the security token service (STS) responsible for creating and delivering the token.
                In case of Azure, it includes the Azure AD tenant where user authentication occurred.
        - 'appid' stands for the application ID of the client utilizing this token. This application can operate either
                autonomously or on behalf of a user. The application ID commonly represents an application object but
                may also denote a service principal object in case of Azure.
        - 'idp' is used to document the identity provider that authenticated the subject of the token.
        - 'oid' is the unchanging identifier for an entity within the identity system.
        - 'sub' identifies the primary entity for the token, such as the user of an app. This value is specific to
                a particular application ID. If a single user logs into two different apps using distinct client IDs,
                these apps will receive different values for the subject claim.
        - 'tid' In case of Azure, this value represents Azure Tenant ID.

        See https://datatracker.ietf.org/doc/html/rfc7519 for specification.
        See https://jwt.ms for debugger.
        """
        try:
            jwt_split = self.access_token.split(".")
            if len(jwt_split) != 3:
                logger.debug(f"Tried to decode access token as JWT, but failed: {len(jwt_split)} components")
                return {}
            payload_with_padding = jwt_split[1] + "=="
            payload_bytes = base64.urlsafe_b64decode(payload_with_padding)
            payload_json = payload_bytes.decode("utf8")
            claims = json.loads(payload_json)
            return claims
        except ValueError as err:
            logger.debug(f"Tried to decode access token as JWT, but failed: {err}")
            return {}


class TokenSource:
    @abstractmethod
    def token(self) -> Token:
        pass


def retrieve_token(
    client_id,
    client_secret,
    token_url,
    params,
    use_params=False,
    use_header=False,
    headers=None,
) -> Token:
    logger.debug(f"Retrieving token for {client_id}")
    if use_params:
        if client_id:
            params["client_id"] = client_id
        if client_secret:
            params["client_secret"] = client_secret
    auth = None
    if use_header:
        auth = requests.auth.HTTPBasicAuth(client_id, client_secret)
    else:
        auth = IgnoreNetrcAuth()
    resp = requests.post(token_url, params, auth=auth, headers=headers)
    if not resp.ok:
        if resp.headers["Content-Type"].startswith("application/json"):
            err = resp.json()
            code = err.get("errorCode", err.get("error", "unknown"))
            summary = err.get("errorSummary", err.get("error_description", "unknown"))
            summary = summary.replace("\r\n", " ")
            raise ValueError(f"{code}: {summary}")
        raise ValueError(resp.content)
    try:
        j = resp.json()
        expires_in = int(j["expires_in"])
        expiry = datetime.now() + timedelta(seconds=expires_in)
        return Token(
            access_token=j["access_token"],
            refresh_token=j.get("refresh_token"),
            token_type=j["token_type"],
            expiry=expiry,
        )
    except Exception as e:
        raise NotImplementedError(f"Not supported yet: {e}")


class _TokenState(Enum):
    """
    Represents the state of a token. Each token can be in one of
    the following three states:
      - FRESH: The token is valid.
      - STALE: The token is valid but will expire soon.
      - EXPIRED: The token has expired and cannot be used.
    """

    FRESH = 1  # The token is valid.
    STALE = 2  # The token is valid but will expire soon.
    EXPIRED = 3  # The token has expired and cannot be used.


class Refreshable(TokenSource):
    """A token source that supports refreshing expired tokens."""

    _EXECUTOR = None
    _EXECUTOR_LOCK = threading.Lock()
    # Default maximum stale duration. Chosen to cover the maximum monthly downtime
    # allowed by a 99.99% uptime SLA (~4.38 minutes) with generous overhead guarantees
    _MAX_STALE_DURATION = timedelta(minutes=20)
    # Backoff time after an async refresh failure before trying another one.
    _ASYNC_REFRESH_RETRY_BACKOFF = timedelta(minutes=1)

    @classmethod
    def _get_executor(cls):
        """Lazy initialization of the ThreadPoolExecutor."""
        if cls._EXECUTOR is None:
            with cls._EXECUTOR_LOCK:
                if cls._EXECUTOR is None:
                    # This thread pool has multiple workers because it is shared by all instances of Refreshable.
                    cls._EXECUTOR = ThreadPoolExecutor(max_workers=10)
        return cls._EXECUTOR

    def __init__(
        self,
        token: Optional[Token] = None,
        disable_async: bool = True,
        stale_duration: Optional[timedelta] = None,
    ):
        # Config properties
        self._use_legacy_stale_duration = stale_duration is not None
        # Only read on the legacy path (when _use_legacy_stale_duration is True).
        self._stale_duration = stale_duration if stale_duration is not None else timedelta(seconds=0)
        self._disable_async = disable_async
        # Lock
        self._lock = threading.Lock()
        # Non Thread safe properties. They should be accessed only when protected by the lock above.
        self._stale_after: Optional[datetime] = None
        self._token_generation: int = 0
        self._update_token(token or Token(""))
        self._is_refreshing = False

    def _now(self) -> datetime:
        """Return the current time, matching the tz-awareness of the cached token."""
        if self._token.expiry:
            return datetime.now(tz=self._token.expiry.tzinfo)
        if self._stale_after:
            return datetime.now(tz=self._stale_after.tzinfo)
        return datetime.now()

    def _update_token(self, token: Token) -> None:
        """Stores the new token and pre-computes the stale threshold.

        The stale period is computed once at token acquisition time as:
            stale_period = min(TTL x 0.5, max_stale_duration)

        This ensures short-lived tokens (e.g. FastPath with 10-minute TTL) get a
        proportionally smaller stale window, while standard OAuth tokens (≥1 hour TTL)
        use the full cap of _MAX_STALE_DURATION.
        """
        self._token = token
        self._token_generation += 1
        self._stale_after = None

        if self._token.expiry:
            if self._use_legacy_stale_duration:
                self._stale_after = self._token.expiry - self._stale_duration
            else:
                ttl = self._token.expiry - self._now()
                stale_duration = max(timedelta(seconds=0), min(ttl // 2, self._MAX_STALE_DURATION))
                self._stale_after = self._token.expiry - stale_duration

    def _handle_failed_async_refresh(self) -> None:
        """Pushes _stale_after forward by the retry backoff, making the token appear fresh temporarily.

        This may set _stale_after past the token's expiry; that is safe because
        _token_state() checks expiry before staleness.
        """
        if self._stale_after:
            self._stale_after = self._now() + self._ASYNC_REFRESH_RETRY_BACKOFF

    # This is the main entry point for the Token. Do not access the token
    # using any of the internal functions.
    def token(self) -> Token:
        """Returns a valid token, blocking if async refresh is disabled."""
        with self._lock:
            if self._disable_async:
                return self._blocking_token()
            return self._async_token()

    def _async_token(self) -> Token:
        """
        Returns a token.
        If the token is stale, triggers an asynchronous refresh.
        If the token is expired, refreshes it synchronously, blocking until the refresh is complete.
        """
        state = self._token_state()
        token = self._token

        if state == _TokenState.FRESH:
            return token
        if state == _TokenState.STALE:
            self._trigger_async_refresh()
            return token
        return self._blocking_token()

    def _token_state(self) -> _TokenState:
        """Returns the current state of the token."""
        if not self._token or not self._token.valid:
            return _TokenState.EXPIRED
        if not self._token.expiry:
            return _TokenState.FRESH

        now = self._now()
        if self._token.expiry < now:
            return _TokenState.EXPIRED
        if self._stale_after and self._stale_after < now:
            return _TokenState.STALE
        return _TokenState.FRESH

    def _blocking_token(self) -> Token:
        """Returns a token, blocking if necessary to refresh it."""
        state = self._token_state()
        self._is_refreshing = False

        # It's possible that the token got refreshed (either by a _blocking_refresh or
        # an _async_refresh call) while this particular call was waiting to acquire
        # the lock. This check avoids refreshing the token again in such cases.
        if state != _TokenState.EXPIRED:
            return self._token

        self._update_token(self.refresh())
        return self._token

    def _trigger_async_refresh(self):
        """Starts an asynchronous refresh if none is in progress."""
        gen_at_submit = self._token_generation

        def _refresh_internal():
            new_token = None
            try:
                new_token = self.refresh()
            except Exception as e:
                # This happens on a thread, so we don't want to propagate the error.
                # Instead, if there is no new_token for any reason, we apply a retry
                # backoff below so the token appears fresh for a short cooldown period.
                logger.warning(f"Tried to refresh token asynchronously, but failed: {e}")

            with self._lock:
                if self._token_generation != gen_at_submit:
                    logger.debug("Async refresh completed but token was already updated; discarding result.")
                elif new_token is not None:
                    self._update_token(new_token)
                else:
                    self._handle_failed_async_refresh()
                self._is_refreshing = False

        # The token may have been refreshed by another thread.
        if self._token_state() == _TokenState.FRESH:
            return
        if not self._is_refreshing:
            self._is_refreshing = True
            Refreshable._get_executor().submit(_refresh_internal)

    @abstractmethod
    def refresh(self) -> Token:
        pass


class _OAuthCallback(BaseHTTPRequestHandler):
    def __init__(self, feedback: list, *args):
        self._feedback = feedback
        super().__init__(*args)

    def log_message(self, fmt: str, *args: Any) -> None:
        logger.debug(fmt, *args)

    def do_GET(self):
        from urllib.parse import parse_qsl

        parts = self.path.split("?")
        if len(parts) != 2:
            self.send_error(400, "Missing Query")
            return

        query = dict(parse_qsl(parts[1]))
        self._feedback.append(query)

        if "error" in query:
            self.send_error(400, query["error"], query.get("error_description"))
            return

        self.send_response(200)
        self.send_header("Content-type", "text/html")
        self.end_headers()
        # TODO: show better message
        self.wfile.write(b"You can close this tab.")


@dataclass
class HostMetadata:
    """Parsed response from the /.well-known/databricks-config discovery endpoint."""

    oidc_endpoint: str
    account_id: Optional[str] = None
    workspace_id: Optional[str] = None
    cloud: Optional[Cloud] = None
    host_type: Optional[str] = None

    @staticmethod
    def from_dict(d: dict) -> "HostMetadata":
        return HostMetadata(
            oidc_endpoint=d.get("oidc_endpoint", ""),
            account_id=d.get("account_id"),
            workspace_id=d.get("workspace_id"),
            cloud=Cloud.parse(d.get("cloud", "")),
            host_type=d.get("host_type"),
        )

    def as_dict(self) -> dict:
        return {
            "oidc_endpoint": self.oidc_endpoint,
            "account_id": self.account_id,
            "workspace_id": self.workspace_id,
            "cloud": self.cloud.value if self.cloud else None,
            "host_type": self.host_type,
        }


def get_host_metadata(host: str, client: _BaseClient = _BaseClient()) -> HostMetadata:
    """
    [Experimental] Fetch the raw Databricks well-known configuration from {host}/.well-known/databricks-config.

    :param host: The Databricks host (workspace or account console).
    :return: Parsed :class:`HostMetadata` as returned by the server.
    """
    host = _fix_host_if_needed(host)
    try:
        resp = client.do("GET", f"{host}/.well-known/databricks-config")
    except Exception as e:
        raise ValueError(f"Failed to fetch host metadata from {host}/.well-known/databricks-config: {e}") from e
    return HostMetadata.from_dict(resp)


def get_endpoints_from_url(url: str, client: _BaseClient = _BaseClient()) -> OidcEndpoints:
    """
    Fetch OIDC endpoints directly from a discovery URL.

    :param url: Full URL to the OIDC discovery document (e.g. the value of discovery_url config).
    :return: Parsed :class:`OidcEndpoints`.
    """
    resp = client.do("GET", url)
    return OidcEndpoints.from_dict(resp)


def get_account_endpoints(host: str, account_id: str, client: _BaseClient = _BaseClient()) -> OidcEndpoints:
    """
    Get the OIDC endpoints for a given account.
    :param host: The Databricks account host.
    :param account_id: The account ID.
    :return: The account's OIDC endpoints.
    """
    host = _fix_host_if_needed(host)
    oidc = f"{host}/oidc/accounts/{account_id}/.well-known/oauth-authorization-server"
    resp = client.do("GET", oidc)
    return OidcEndpoints.from_dict(resp)


def get_workspace_endpoints(host: str, client: _BaseClient = _BaseClient()) -> OidcEndpoints:
    """
    Get the OIDC endpoints for a given workspace.
    :param host: The Databricks workspace host.
    :return: The workspace's OIDC endpoints.
    """
    host = _fix_host_if_needed(host)
    oidc = f"{host}/oidc/.well-known/oauth-authorization-server"
    resp = client.do("GET", oidc)
    return OidcEndpoints.from_dict(resp)


def get_unified_endpoints(host: str, account_id: str, client: _BaseClient = _BaseClient()) -> OidcEndpoints:
    """
    Get the OIDC endpoints for a unified host.
    :param host: The Databricks unified host.
    :param account_id: The account ID.
    :return: The OIDC endpoints for the unified host.
    """
    host = _fix_host_if_needed(host)
    oidc = f"{host}/oidc/accounts/{account_id}/.well-known/oauth-authorization-server"
    resp = client.do("GET", oidc)
    return OidcEndpoints.from_dict(resp)


def get_azure_entra_id_workspace_endpoints(
    host: str,
) -> Optional[OidcEndpoints]:
    """
    Get the Azure Entra ID endpoints for a given workspace. Can only be used when authenticating to Azure Databricks
    using an application registered in Azure Entra ID.
    :param host: The Databricks workspace host.
    :return: The OIDC endpoints for the workspace's Azure Entra ID tenant.
    """
    # In Azure, this workspace endpoint redirects to the Entra ID authorization endpoint
    host = _fix_host_if_needed(host)
    res = requests.get(f"{host}/oidc/oauth2/v2.0/authorize", allow_redirects=False)
    real_auth_url = res.headers.get("location")
    if not real_auth_url:
        return None
    return OidcEndpoints(
        authorization_endpoint=real_auth_url,
        token_endpoint=real_auth_url.replace("/authorize", "/token"),
    )


class SessionCredentials(Refreshable):
    def __init__(
        self,
        token: Token,
        token_endpoint: str,
        client_id: str,
        client_secret: str = None,
        redirect_url: str = None,
        disable_async: bool = True,
    ):
        self._token_endpoint = token_endpoint
        self._client_id = client_id
        self._client_secret = client_secret
        self._redirect_url = redirect_url
        super().__init__(
            token=token,
            disable_async=disable_async,
        )

    def as_dict(self) -> dict:
        return {"token": self.token().as_dict()}

    @staticmethod
    def from_dict(
        raw: dict,
        token_endpoint: str,
        client_id: str,
        client_secret: str = None,
        redirect_url: str = None,
    ) -> "SessionCredentials":
        return SessionCredentials(
            token=Token.from_dict(raw["token"]),
            token_endpoint=token_endpoint,
            client_id=client_id,
            client_secret=client_secret,
            redirect_url=redirect_url,
        )

    def auth_type(self):
        """Implementing CredentialsProvider protocol"""
        # TODO: distinguish between Databricks IDP and Azure AD
        return "oauth"

    def __call__(self, *args, **kwargs):
        """Implementing CredentialsProvider protocol"""

        def inner() -> Dict[str, str]:
            return {"Authorization": f"Bearer {self.token().access_token}"}

        return inner

    def refresh(self) -> Token:
        refresh_token = self._token.refresh_token
        if not refresh_token:
            raise ValueError("oauth2: token expired and refresh token is not set")
        params = {
            "grant_type": "refresh_token",
            "refresh_token": refresh_token,
        }
        headers = {}
        if "microsoft" in self._token_endpoint:
            # Tokens issued for the 'Single-Page Application' client-type may
            # only be redeemed via cross-origin requests
            headers = {"Origin": self._redirect_url}
        return retrieve_token(
            client_id=self._client_id,
            client_secret=self._client_secret,
            token_url=self._token_endpoint,
            params=params,
            use_params=True,
            headers=headers,
        )


class Consent:
    def __init__(
        self,
        state: str,
        verifier: str,
        authorization_url: str,
        redirect_url: str,
        token_endpoint: str,
        client_id: str,
        client_secret: str = None,
    ) -> None:
        self._verifier = verifier
        self._state = state
        self._authorization_url = authorization_url
        self._redirect_url = redirect_url
        self._token_endpoint = token_endpoint
        self._client_id = client_id
        self._client_secret = client_secret

    def as_dict(self) -> dict:
        return {
            "state": self._state,
            "verifier": self._verifier,
            "authorization_url": self._authorization_url,
            "redirect_url": self._redirect_url,
            "token_endpoint": self._token_endpoint,
            "client_id": self._client_id,
        }

    @property
    def authorization_url(self) -> str:
        return self._authorization_url

    @staticmethod
    def from_dict(raw: dict, client_secret: str = None) -> "Consent":
        return Consent(
            raw["state"],
            raw["verifier"],
            authorization_url=raw["authorization_url"],
            redirect_url=raw["redirect_url"],
            token_endpoint=raw["token_endpoint"],
            client_id=raw["client_id"],
            client_secret=client_secret,
        )

    def launch_external_browser(self) -> SessionCredentials:
        redirect_url = urllib.parse.urlparse(self._redirect_url)
        if redirect_url.hostname not in ("localhost", "127.0.0.1"):
            raise ValueError(f"cannot listen on {redirect_url.hostname}")
        feedback = []
        logger.info(f"Opening {self._authorization_url} in a browser")
        webbrowser.open_new(self._authorization_url)
        port = redirect_url.port
        handler_factory = functools.partial(_OAuthCallback, feedback)
        with HTTPServer(("localhost", port), handler_factory) as httpd:
            logger.info(f"Waiting for redirect to http://localhost:{port}")
            httpd.handle_request()
        if not feedback:
            raise ValueError("No data received in callback")
        query = feedback.pop()
        return self.exchange_callback_parameters(query)

    def exchange_callback_parameters(self, query: Dict[str, str]) -> SessionCredentials:
        if "error" in query:
            raise ValueError("{error}: {error_description}".format(**query))
        if "code" not in query or "state" not in query:
            raise ValueError("No code returned in callback")
        return self.exchange(query["code"], query["state"])

    def exchange(self, code: str, state: str) -> SessionCredentials:
        if self._state != state:
            raise ValueError("state mismatch")
        params = {
            "redirect_uri": self._redirect_url,
            "grant_type": "authorization_code",
            "code_verifier": self._verifier,
            "code": code,
        }
        headers = {}
        while True:
            try:
                token = retrieve_token(
                    client_id=self._client_id,
                    client_secret=self._client_secret,
                    token_url=self._token_endpoint,
                    params=params,
                    headers=headers,
                    use_params=True,
                )
                return SessionCredentials(
                    token,
                    self._token_endpoint,
                    self._client_id,
                    self._client_secret,
                    self._redirect_url,
                )
            except ValueError as e:
                if NO_ORIGIN_FOR_SPA_CLIENT_ERROR in str(e):
                    # Retry in cases of 'Single-Page Application' client-type with
                    # 'Origin' header equal to client's redirect URL.
                    headers["Origin"] = self._redirect_url
                    msg = f"Retrying OAuth token exchange with {self._redirect_url} origin"
                    logger.debug(msg)
                    continue
                raise e


class OAuthClient:
    """Enables 3-legged OAuth2 flow with PKCE

    For a regular web app running on a server, it's recommended to use
    the Authorization Code Flow to obtain an Access Token and a Refresh
    Token. This method is considered safe because the Access Token is
    transmitted directly to the server hosting the app, without passing
    through the user's web browser and risking exposure.

    To enhance the security of the Authorization Code Flow, the PKCE
    (Proof Key for Code Exchange) mechanism can be employed. With PKCE,
    the calling application generates a secret called the Code Verifier,
    which is verified by the authorization server. The app also creates
    a transform value of the Code Verifier, called the Code Challenge,
    and sends it over HTTPS to obtain an Authorization Code.
    By intercepting the Authorization Code, a malicious attacker cannot
    exchange it for a token without possessing the Code Verifier.
    """

    def __init__(
        self,
        oidc_endpoints: OidcEndpoints,
        redirect_url: str,
        client_id: str,
        scopes: List[str] = None,
        client_secret: str = None,
    ):
        if not scopes:
            # Default for direct OAuthClient users (e.g., via from_host()).
            # When used via credentials_provider.external_browser(), scopes are always
            # passed explicitly from Config.get_scopes(), with offline_access handling
            # controlled by the disable_oauth_refresh_token flag.
            scopes = ["all-apis", "offline_access"]

        self.redirect_url = redirect_url
        self._client_id = client_id
        self._client_secret = client_secret
        self._oidc_endpoints = oidc_endpoints
        self._scopes = scopes

    @staticmethod
    def from_host(
        host: str,
        client_id: str,
        redirect_url: str,
        *,
        scopes: List[str] = None,
        client_secret: str = None,
    ) -> "OAuthClient":
        from .core import Config
        from .credentials_provider import credentials_strategy

        @credentials_strategy("noop", [])
        def noop_credentials(_: any):
            return lambda: {}

        config = Config(host=host, credentials_strategy=noop_credentials)
        oidc = config.databricks_oidc_endpoints
        if not oidc:
            raise ValueError(f"{host} does not support OAuth")
        return OAuthClient(oidc, redirect_url, client_id, scopes, client_secret)

    def initiate_consent(self) -> Consent:
        state = secrets.token_urlsafe(16)

        # token_urlsafe() already returns base64-encoded string
        verifier = secrets.token_urlsafe(32)
        digest = hashlib.sha256(verifier.encode("UTF-8")).digest()
        challenge = base64.urlsafe_b64encode(digest).decode("UTF-8").replace("=", "")

        params = {
            "response_type": "code",
            "client_id": self._client_id,
            "redirect_uri": self.redirect_url,
            "scope": " ".join(self._scopes),
            "state": state,
            "code_challenge": challenge,
            "code_challenge_method": "S256",
        }
        auth_url = f"{self._oidc_endpoints.authorization_endpoint}?{urllib.parse.urlencode(params)}"
        return Consent(
            state,
            verifier,
            authorization_url=auth_url,
            redirect_url=self.redirect_url,
            token_endpoint=self._oidc_endpoints.token_endpoint,
            client_id=self._client_id,
            client_secret=self._client_secret,
        )

    def __repr__(self) -> str:
        return f"<OAuthClient client_id={self._client_id} token_url={self._oidc_endpoints.token_endpoint} auth_url={self._oidc_endpoints.authorization_endpoint}>"


@dataclass
class ClientCredentials(Refreshable):
    """Enables client credentials 2-legged OAuth2 flow

    When it comes to authorizing machine-to-machine interactions,
    the need for end-user authorization is eliminated because the SDK
    functions as both the Resource Owner and Client. Typical example is
    the CI/CD process or any other automated job. In this scenario,
    the background job uses the Client ID and Client Secret to obtain
    an Access Token from the Authorization Server.
    """

    client_id: str
    client_secret: str
    token_url: str
    endpoint_params: dict = None
    scopes: str = None
    use_params: bool = False
    use_header: bool = False
    disable_async: bool = True
    authorization_details: str = None

    def __post_init__(self):
        super().__init__(disable_async=self.disable_async)

    def refresh(self) -> Token:
        params = {"grant_type": "client_credentials"}
        if self.scopes:
            params["scope"] = self.scopes
        if self.authorization_details:
            params["authorization_details"] = self.authorization_details
        if self.endpoint_params:
            for k, v in self.endpoint_params.items():
                params[k] = v
        return retrieve_token(
            self.client_id,
            self.client_secret,
            self.token_url,
            params,
            use_params=self.use_params,
            use_header=self.use_header,
        )


@dataclass
class PATOAuthTokenExchange(Refreshable):
    """Performs OAuth token exchange using a Personal Access Token (PAT) as the subject token.

    This class implements the OAuth 2.0 Token Exchange flow (RFC 8693) to exchange a Databricks
    Internal PAT Token for an access token with specific scopes and authorization details.

    Args:
        get_original_token: A callable that returns the PAT to be exchanged. This is a callable
            rather than a string value to ensure that a fresh Internal PAT Token is retrieved
            at the time of refresh.
        host: The Databricks workspace URL (e.g., "https://my-workspace.cloud.databricks.com").
        scopes: Space-delimited string of OAuth scopes to request (e.g., "all-apis offline_access").
        authorization_details: Optional JSON string containing authorization details as defined in
            AuthorizationDetail class above.
        disable_async: Whether to disable asynchronous token refresh. Defaults to True.
    """

    get_original_token: Callable[[], Optional[str]]
    host: str
    scopes: str
    authorization_details: str = None
    disable_async: bool = True

    def __post_init__(self):
        super().__init__(disable_async=self.disable_async)

    def refresh(self) -> Token:
        token_exchange_url = f"{self.host}/oidc/v1/token"
        params = {
            "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
            "subject_token": self.get_original_token(),
            "subject_token_type": "urn:databricks:params:oauth:token-type:personal-access-token",
            "requested_token_type": "urn:ietf:params:oauth:token-type:access_token",
            "scope": self.scopes,
        }
        if self.authorization_details:
            params["authorization_details"] = self.authorization_details

        resp = requests.post(token_exchange_url, params)
        if not resp.ok:
            if resp.headers["Content-Type"].startswith("application/json"):
                err = resp.json()
                code = err.get("errorCode", err.get("error", "unknown"))
                summary = err.get("errorSummary", err.get("error_description", "unknown"))
                summary = summary.replace("\r\n", " ")
                raise ValueError(f"{code}: {summary}")
            raise ValueError(resp.content)
        try:
            j = resp.json()
            expires_in = int(j["expires_in"])
            expiry = datetime.now() + timedelta(seconds=expires_in)
            return Token(
                access_token=j["access_token"],
                expiry=expiry,
                token_type=j["token_type"],
            )
        except Exception as e:
            raise ValueError(f"Failed to exchange PAT for OAuth token: {e}")


class TokenCache:
    BASE_PATH = "~/.config/databricks-sdk-py/oauth"

    def __init__(
        self,
        host: str,
        oidc_endpoints: OidcEndpoints,
        client_id: str,
        redirect_url: Optional[str] = None,
        client_secret: Optional[str] = None,
        scopes: Optional[List[str]] = None,
        profile: Optional[str] = None,
    ) -> None:
        self._host = host
        self._client_id = client_id
        self._oidc_endpoints = oidc_endpoints
        self._redirect_url = redirect_url
        self._client_secret = client_secret
        self._scopes = scopes or []
        self._profile = profile

    @property
    def filename(self) -> str:
        # Include host, client_id, scopes, and profile in the cache filename to make it unique.
        # JSON serialization ensures values are properly escaped and separated.
        key = {
            "host": self._host,
            "client_id": self._client_id,
            "scopes": self._scopes,
            "profile": self._profile or "",
        }
        h = hashlib.sha256(json.dumps(key, sort_keys=True).encode("utf-8"))
        return os.path.expanduser(os.path.join(self.__class__.BASE_PATH, h.hexdigest() + ".json"))

    def load(self) -> Optional[SessionCredentials]:
        """
        Load credentials from cache file. Return None if the cache file does not exist or is invalid.
        """
        if not os.path.exists(self.filename):
            return None

        try:
            with open(self.filename, "r") as f:
                raw = json.load(f)
                return SessionCredentials.from_dict(
                    raw,
                    token_endpoint=self._oidc_endpoints.token_endpoint,
                    client_id=self._client_id,
                    client_secret=self._client_secret,
                    redirect_url=self._redirect_url,
                )
        except Exception:
            return None

    def save(self, credentials: SessionCredentials) -> None:
        """
        Save credentials to cache file.
        """
        os.makedirs(os.path.dirname(self.filename), exist_ok=True)
        with open(self.filename, "w") as f:
            json.dump(credentials.as_dict(), f)
        os.chmod(self.filename, 0o600)