_token_client.py

import asyncio
import base64
import enum
import typing
import urllib.parse
from datetime import datetime, timedelta

import httpx
import pydantic

from flyte._logging import logger
from flyte.remote._client.auth.errors import AuthenticationError, AuthenticationPending

utf_8 = "utf-8"

# Errors that Token endpoint will return
error_slow_down = "slow_down"
error_auth_pending = "authorization_pending"


# Grant Types
class GrantType(str, enum.Enum):
    CLIENT_CREDS = "client_credentials"
    DEVICE_CODE = "urn:ietf:params:oauth:grant-type:device_code"
    REFRESH_TOKEN = "refresh_token"


class DeviceCodeResponse(pydantic.BaseModel):
    """
    Response from device auth flow endpoint
    {
        'device_code': 'code',
         'user_code': 'BNDJJFXL',
         'verification_uri': 'url',
         'expires_in': 600,
         'interval': 5
    }

    Attributes:
        device_code (str): The device verification code.
        user_code (str): The user-facing code that should be entered on the verification page.
        verification_uri (str): The URL where the user should enter the user_code.
        expires_in (int): The lifetime in seconds of the device code and user code.
        interval (int): The minimum amount of time in seconds to wait between polling requests.
    """

    device_code: str
    user_code: str
    verification_uri: str
    expires_in: int
    interval: int

    @classmethod
    def from_json_response(cls, j: typing.Dict) -> "DeviceCodeResponse":
        """
        Create a DeviceCodeResponse instance from a JSON response dictionary.

        Args:
            j: The JSON response dictionary containing device code information

        Returns:
            A new instance with values from the JSON response
        """
        return cls(
            device_code=j["device_code"],
            user_code=j["user_code"],
            verification_uri=j["verification_uri"],
            expires_in=j["expires_in"],
            interval=j["interval"],
        )


def _body_snippet(response: httpx.Response, limit: int = 200) -> str:
    """Describe a response body compactly enough to put in an error message."""
    content_type = response.headers.get("content-type", "unknown")
    text = " ".join(response.text.split())
    if len(text) > limit:
        text = text[:limit] + "..."
    return f"content-type: {content_type}, body: {text!r}" if text else f"content-type: {content_type}, empty body"


def _json_object_or_none(response: httpx.Response) -> typing.Optional[typing.Dict]:
    """Parse an IDP response body as a JSON object, or None when it is anything else.

    The OAuth endpoints are specified to answer in JSON, but what actually reaches the SDK is
    whatever sits between it and the IDP: a load balancer's HTML 502 page, a proxy's plain-text
    "Internal Server Error", an SSO interstitial. That is a deployment problem, and the branches
    below already know how to report it -- they only have to survive reading the body first.

    Returns None rather than raising, and only for a JSON *object*: a body that parses to a bare
    string would still answer `"error" in j` by substring, which is not the membership test the
    caller means.
    """
    try:
        parsed = response.json()
    except ValueError:
        return None
    return parsed if isinstance(parsed, dict) else None


def get_basic_authorization_header(client_id: str, client_secret: str) -> str:
    """
    This function transforms the client id and the client secret into a header that conforms with http basic auth.
    It joins the id and the secret with a : then base64 encodes it, then adds the appropriate text. Secrets are
    first URL encoded to escape illegal characters.

    Args:
        client_id: The client ID for authentication
        client_secret: The client secret for authentication

    Returns:
        str
    """
    encoded = urllib.parse.quote_plus(client_secret)
    concatenated = "{}:{}".format(client_id, encoded)
    return "Basic {}".format(base64.b64encode(concatenated.encode(utf_8)).decode(utf_8))


async def get_token(
    token_endpoint: str,
    http_session: httpx.AsyncClient,
    scopes: typing.Optional[typing.List[str]] = None,
    authorization_header: typing.Optional[str] = None,
    client_id: typing.Optional[str] = None,
    device_code: typing.Optional[str] = None,
    audience: typing.Optional[str] = None,
    grant_type: GrantType = GrantType.CLIENT_CREDS,
    http_proxy_url: typing.Optional[str] = None,
    verify: typing.Optional[typing.Union[bool, str]] = None,
    refresh_token: typing.Optional[str] = None,
) -> typing.Tuple[str, str | None, int]:
    """
    Retrieves an access token from the specified token endpoint.

    Args:
        token_endpoint: The endpoint URL for token retrieval
        http_session: HTTP session to use for requests
        scopes: Optional list of scopes to request during authentication
        authorization_header: Optional authorization header value
        client_id: Optional client ID for authentication
        device_code: Optional device code for device flow authentication
        audience: Optional audience for the token
        grant_type: The grant type to use (default: CLIENT_CREDS)
        http_proxy_url: Optional HTTP proxy URL
        verify: Whether to verify SSL certificates (bool or path to cert)
        refresh_token: Optional refresh token for token refresh

    Returns:
        A tuple containing (access_token, refresh_token, expires_in)

    Raises:
        AuthenticationPending: When authentication is still pending (for device code flow).
        AuthenticationError: When authentication fails for any reason.
    """
    headers = {
        "Cache-Control": "no-cache",
        "Accept": "application/json",
        "Content-Type": "application/x-www-form-urlencoded",
    }
    if authorization_header:
        headers["Authorization"] = authorization_header
    body = {
        "grant_type": grant_type.value,
    }
    if client_id:
        body["client_id"] = client_id
    if device_code:
        body["device_code"] = device_code
    if scopes is not None:
        body["scope"] = " ".join(s.strip("' ") for s in scopes).strip("[]'")
    if audience:
        body["audience"] = audience
    if refresh_token:
        body["refresh_token"] = refresh_token

    response = await http_session.post(token_endpoint, data=body, headers=headers)

    if not response.is_success:
        j = _json_object_or_none(response)
        if j is not None and "error" in j:
            err = j["error"]
            if err == error_auth_pending or err == error_slow_down:
                raise AuthenticationPending(f"Token not yet available, try again in some time {err}")
        logger.error("Status Code ({}) received from IDP: {}".format(response.status_code, response.text))
        raise AuthenticationError("Status Code ({}) received from IDP: {}".format(response.status_code, response.text))

    j = _json_object_or_none(response)
    if j is None or "access_token" not in j:
        # A 2xx that is not a usable token response: an authenticating proxy answering with its
        # own login page, or an endpoint that is not the IDP's at all. Saying so beats the
        # KeyError/JSONDecodeError that used to escape from here.
        raise AuthenticationError(
            f"Token endpoint {token_endpoint} returned {response.status_code} but not an access "
            f"token ({_body_snippet(response)}). Check that the endpoint in your config points at "
            f"the identity provider and that nothing is intercepting the request."
        )

    new_refresh_token = None
    if "refresh_token" in j:
        new_refresh_token = j["refresh_token"]
    else:
        logger.info("No refresh token received, this is expected for client credentials flow")

    return j["access_token"], new_refresh_token, j["expires_in"]


async def get_device_code(
    device_auth_endpoint: str,
    client_id: str,
    http_session: httpx.AsyncClient,
    *,
    audience: typing.Optional[str] = None,
    scopes: typing.Optional[typing.List[str]] = None,
) -> DeviceCodeResponse:
    """
    Retrieves the device authentication code that can be used to authenticate the request using a browser on a
    separate device.

    Args:
        device_auth_endpoint: The URL of the device authorization endpoint
        client_id: The client ID to use for authentication
        audience: The audience value to request
        scopes: List of scopes to request
        http_proxy_url: HTTP proxy URL if needed
        verify: SSL verification mode
        http_session: An existing HTTP client session

    Returns:
        An object containing the device code and related information

    Raises:
        AuthenticationError: When device code retrieval fails
    """
    _scope = " ".join(s.strip("' ") for s in scopes).strip("[]'") if scopes is not None else ""
    payload = {"client_id": client_id, "scope": _scope, "audience": audience}
    resp = await http_session.post(device_auth_endpoint, data=payload)
    if not resp.is_success:
        raise AuthenticationError(
            f"Unable to retrieve Device Authentication Code for {payload},"
            f" Status Code {resp.status_code} Reason {_body_snippet(resp)}"
        )
    j = _json_object_or_none(resp)
    if j is None:
        raise AuthenticationError(
            f"Device authorization endpoint {device_auth_endpoint} returned {resp.status_code} "
            f"with a body that is not a JSON object ({_body_snippet(resp)}). Check that the "
            f"endpoint in your config points at the identity provider."
        )
    return DeviceCodeResponse.from_json_response(j)


async def poll_token_endpoint(
    resp: DeviceCodeResponse,
    *,
    token_endpoint: str,
    client_id: str,
    http_session: httpx.AsyncClient,
    audience: typing.Optional[str] = None,
    scopes: typing.Optional[typing.List[str]] = None,
    http_proxy_url: typing.Optional[str] = None,
    verify: typing.Optional[typing.Union[bool, str]] = None,
) -> typing.Tuple[str, str | None, int]:
    """
    Polls the token endpoint until authentication is complete or times out.

    This function repeatedly calls the token endpoint at the specified interval until either:
    1. Authentication is successful and a token is returned
    2. The device code expires (as specified in the DeviceCodeResponse)

    Args:
        resp: The device code response from a previous call to get_device_code
        token_endpoint: The URL of the token endpoint
        client_id: The client ID to use for authentication
        audience: The audience value to request
        scopes: Space-separated list of scopes to request
        http_proxy_url: HTTP proxy URL if needed
        verify: SSL verification mode

    Returns:
        A tuple containing (access_token, refresh_token, expires_in)

    Raises:
        AuthenticationError: When authentication fails or times out
    """
    tick = datetime.now()
    interval = timedelta(seconds=resp.interval)
    end_time = tick + timedelta(seconds=resp.expires_in)
    while tick < end_time:
        try:
            access_token, refresh_token, expires_in = await get_token(
                token_endpoint,
                grant_type=GrantType.DEVICE_CODE,
                client_id=client_id,
                audience=audience,
                scopes=scopes,
                device_code=resp.device_code,
                http_proxy_url=http_proxy_url,
                verify=verify,
                http_session=http_session,
            )
            logger.debug(f"Authentication successful, access token received, expires in {expires_in} seconds")
            return access_token, refresh_token, expires_in
        except AuthenticationPending:
            ...
        except Exception as e:
            logger.warning(f"Authentication failed, reason {e}")
            raise e
        logger.debug(f"Authentication pending, ..., waiting for {resp.interval} seconds")
        await asyncio.sleep(interval.total_seconds())
        tick = tick + interval
    raise AuthenticationError("Authentication failed!")