flyte 0.0.1b0__py3-none-any.whl

union/remote/_client/auth/_authenticators/device_code.py

@@ -0,0 +1,120 @@
+import click
+
+from union._logging import logger
+from union.remote._client.auth import _token_client as token_client
+from union.remote._client.auth._authenticators.base import Authenticator
+from union.remote._client.auth._keyring import Credentials
+from union.remote._client.auth.errors import AuthenticationError, AuthenticationPending
+
+
+class DeviceCodeAuthenticator(Authenticator):
+    """
+    This Authenticator implements the Device Code authorization flow useful for headless user authentication.
+
+    Examples described
+    - https://developer.okta.com/docs/guides/device-authorization-grant/main/
+    - https://auth0.com/docs/get-started/authentication-and-authorization-flow/device-authorization-flow#device-flow
+    """
+
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        """
+        Initialize the device code authenticator.
+
+        :param kwargs: Keyword arguments passed to the base Authenticator
+
+        **Keyword Arguments passed to base Authenticator**:
+        :param endpoint: The endpoint URL for authentication (required)
+        :param cfg_store: Optional client configuration store for retrieving remote configuration
+        :param client_config: Optional client configuration containing authentication settings
+        :param credentials: Optional credentials to use for authentication
+        :param http_session: Optional HTTP session to use for requests
+        :param http_proxy_url: Optional HTTP proxy URL
+        :param verify: Whether to verify SSL certificates (default: True)
+        :param ca_cert_path: Optional path to CA certificate file
+        :param client_id: Client ID for authentication
+        :param scopes: List of scopes to request during authentication
+        :param audience: Audience for the token
+        :param device_authorization_endpoint: Endpoint for device authorization
+        """
+
+        super().__init__(
+            **kwargs,
+        )
+
+    async def _do_refresh_credentials(self) -> Credentials:
+        """
+        Refreshes the authentication credentials using device code flow.
+
+        First attempts to refresh using a refresh token if available.
+        If that fails, falls back to the full device code authorization flow.
+        """
+        cfg = await self._resolve_config()
+
+        # These always come from the public client config
+        if cfg.device_authorization_endpoint is None:
+            raise AuthenticationError(
+                "Device Authentication is not available on the Flyte backend / authentication server"
+            )
+
+        if self._creds and self._creds.refresh_token:
+            """We have an refresh token so lets try to refresh it"""
+            try:
+                access_token, refresh_token, expires_in = await token_client.get_token(
+                    token_endpoint=cfg.token_endpoint,
+                    client_id=cfg.client_id,
+                    audience=cfg.audience,
+                    scopes=cfg.scopes,
+                    http_proxy_url=self._http_proxy_url,
+                    verify=self._verify,
+                    grant_type=token_client.GrantType.REFRESH_TOKEN,
+                    refresh_token=self._creds.refresh_token,
+                    http_session=self._http_session,
+                )
+
+                return Credentials(
+                    access_token=access_token,
+                    refresh_token=refresh_token,
+                    expires_in=expires_in,
+                    for_endpoint=self._endpoint,
+                )
+            except (AuthenticationError, AuthenticationPending):
+                logger.warning("Failed to refresh token. Kicking off a full authorization flow.")
+
+        """Fall back to device flow"""
+        resp = await token_client.get_device_code(
+            cfg.device_authorization_endpoint,
+            cfg.client_id,
+            audience=cfg.audience,
+            scopes=cfg.scopes,
+            http_proxy_url=self._http_proxy_url,
+            verify=self._verify,
+            http_session=self._http_session,
+        )
+
+        full_uri = f"{resp.verification_uri}?user_code={resp.user_code}"
+        text = (
+            f"To Authenticate, navigate in a browser to the following URL: "
+            f"{click.style(full_uri, fg='blue', underline=True)}"
+        )
+        click.secho(text)
+        try:
+            token, refresh_token, expires_in = await token_client.poll_token_endpoint(
+                resp,
+                token_endpoint=cfg.token_endpoint,
+                client_id=cfg.client_id,
+                audience=cfg.audience,
+                scopes=cfg.scopes,
+                http_proxy_url=self._http_proxy_url,
+                verify=self._verify,
+                http_session=self._http_session,
+            )
+
+            return Credentials(
+                access_token=token, refresh_token=refresh_token, expires_in=expires_in, for_endpoint=self._endpoint
+            )
+
+        except Exception:
+            raise

union/remote/_client/auth/_authenticators/external_command.py

@@ -0,0 +1,77 @@
+import asyncio
+import typing
+
+from union._logging import logger
+from union.remote._client.auth._authenticators.base import Authenticator
+from union.remote._client.auth._keyring import Credentials
+from union.remote._client.auth.errors import AuthenticationError
+
+
+class AsyncCommandAuthenticator(Authenticator):
+    """
+    This Authenticator retrieves access_token using the provided command
+    """
+
+    def __init__(self, command: typing.List[str], **kwargs):
+        """
+        Initialize the command authenticator.
+
+        :param command: List of command strings to execute for token retrieval
+        :param kwargs: Additional keyword arguments passed to the base Authenticator
+
+        **Keyword Arguments passed to base Authenticator**:
+        :param endpoint: The endpoint URL for authentication
+        :param cfg_store: Optional client configuration store for retrieving remote configuration
+        :param client_config: Optional client configuration containing authentication settings
+        :param credentials: Optional credentials to use for authentication
+        :param http_session: Optional HTTP session to use for requests
+        :param http_proxy_url: Optional HTTP proxy URL
+        :param verify: Whether to verify SSL certificates (default: True)
+        :param ca_cert_path: Optional path to CA certificate file
+
+        **Additional Keyword Arguments**:
+        :param header_key: Header key to use for authentication (defaults to "authorization")
+        :param proxy_env: Environment variables for the command execution
+        :param proxy_timeout: Timeout for command execution
+        """
+        self._cmd = command
+        if not self._cmd:
+            raise AuthenticationError("Command cannot be empty for command authenticator")
+        super().__init__(**kwargs)
+
+    async def _do_refresh_credentials(self) -> Credentials:
+        """
+        Refreshes the authentication credentials by executing an external command.
+
+        This function is used when the configuration value for AUTH_MODE is set to 'external_process'.
+        It reads an id token generated by an external process started by running the 'command'.
+        Uses asyncio.create_subprocess_exec for non-blocking operation.
+
+        The command is executed with stdout and stderr captured, and the stdout output is used
+        as the access token for authentication.
+
+        :raises AuthenticationError: If the command fails to execute or returns a non-zero exit code
+        """
+        cmd_joined = " ".join(self._cmd)
+        logger.debug("Starting external process to generate id token. Command `{}`".format(" ".join(cmd_joined)))
+        try:
+            # Use asyncio subprocess for non-blocking operation
+            process = await asyncio.create_subprocess_exec(
+                *self._cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+            )
+            stdout, stderr = await process.communicate()
+
+            if process.returncode != 0:
+                logger.error(f"Failed to generate token from command `{cmd_joined}`. Error: {stderr.decode()}")
+                raise AuthenticationError(
+                    f"Failed to refresh token with command `{cmd_joined}`."
+                    f" Please execute this command in your terminal to debug."
+                )
+
+            return Credentials(for_endpoint=self._endpoint, access_token=stdout.decode().strip())
+        except Exception as e:
+            logger.error(f"Failed to generate token from command `{cmd_joined}`. Error: {e!s}")
+            raise AuthenticationError(
+                f"Failed to refresh token with command `{cmd_joined}`."
+                f" Please execute this command in your terminal to debug."
+            )

union/remote/_client/auth/_authenticators/factory.py

@@ -0,0 +1,200 @@
+import typing
+
+import grpc.aio
+
+from union.remote._client.auth._authenticators.base import Authenticator
+from union.remote._client.auth._authenticators.external_command import (
+    AsyncCommandAuthenticator,
+)
+from union.remote._client.auth._client_config import AuthType, ClientConfigStore, RemoteClientConfigStore
+
+
+def create_auth_interceptors(
+    endpoint: str, in_channel: grpc.aio.Channel, **kwargs
+) -> typing.List[grpc.aio.ClientInterceptor]:
+    """
+    Async version of upgrade_channel_to_authenticated.
+    Given a grpc.Channel, preferably a secure channel, it returns a list of interceptors to
+    perform an Oauth2.0 Auth flow for all RPC call types.
+
+    :param endpoint: The endpoint URL for authentication
+    :param in_channel: grpc.Channel Precreated channel
+    :param kwargs: Additional arguments passed to the authenticator, including:
+        - insecure: Whether to use an insecure channel
+        - insecure_skip_verify: Whether to skip SSL certificate verification
+        - ca_cert_file_path: Path to CA certificate file for SSL verification
+        - auth_type: The authentication type to use ("Pkce", "ClientSecret", "ExternalCommand", "DeviceFlow")
+        - command: Command to execute for ExternalCommand authentication
+        - client_id: Client ID for ClientSecret authentication
+        - client_secret: Client secret for ClientSecret authentication
+        - scopes: List of scopes to request during authentication
+        - audience: Audience for the token
+        - http_proxy_url: HTTP proxy URL
+        - http_session: httpx.AsyncClient session
+        - verify: Whether to verify SSL certificates
+        - ca_cert_path: Optional path to CA certificate file
+        - header_key: Header key to use for authentication
+        - proxy_env: Environment variables for proxy command
+        - proxy_timeout: Timeout for proxy command execution
+        - redirect_uri: OAuth2 redirect URI for PKCE authentication
+        - add_request_auth_code_params_to_request_access_token_params: Whether to add auth code params to token request
+        - request_auth_code_params: Parameters to add to login URI opened in browser
+        - request_access_token_params: Parameters to add when exchanging auth code for access token
+        - refresh_access_token_params: Parameters to add when refreshing access token
+    :return: List of gRPC interceptors for different call types
+    """
+    from union.remote._client.auth._grpc_utils.auth_interceptor import (
+        AuthStreamStreamInterceptor,
+        AuthStreamUnaryInterceptor,
+        AuthUnaryStreamInterceptor,
+        AuthUnaryUnaryInterceptor,
+    )
+
+    def authenticator_factory() -> Authenticator:
+        return get_async_authenticator(endpoint=endpoint, cfg_store=RemoteClientConfigStore(in_channel), **kwargs)
+
+    return [
+        AuthUnaryUnaryInterceptor(authenticator_factory),
+        AuthUnaryStreamInterceptor(authenticator_factory),
+        AuthStreamUnaryInterceptor(authenticator_factory),
+        AuthStreamStreamInterceptor(authenticator_factory),
+    ]
+
+
+def create_proxy_auth_interceptors(
+    endpoint: str, proxy_command: typing.Optional[typing.List[str]] = None, **kwargs
+) -> typing.List[grpc.aio.ClientInterceptor]:
+    """
+    Async version of upgrade_channel_to_proxy_authenticated.
+    If activated in the platform config, given a grpc.Channel, preferably a secure channel, it returns a list of
+    interceptors to perform authentication with a proxy in front of Flyte for all RPC call types.
+
+    :param endpoint: The endpoint URL for authentication
+    :param proxy_command: Command to execute to get proxy authentication token
+    :param kwargs: Additional arguments passed to the authenticator, including:
+        - proxy_env: Environment variables for the proxy command
+        - proxy_timeout: Timeout for the proxy command
+        - header_key: Header key to use for authentication (defaults to "proxy-authorization")
+        - http_session: httpx.AsyncClient session to use for requests
+        - verify: Whether to verify SSL certificates
+        - ca_cert_path: Optional path to CA certificate file
+    :return: List of gRPC interceptors for different call types
+    """
+    if proxy_command:
+        from union.remote._client.auth._grpc_utils.auth_interceptor import (
+            AuthStreamStreamInterceptor,
+            AuthStreamUnaryInterceptor,
+            AuthUnaryStreamInterceptor,
+            AuthUnaryUnaryInterceptor,
+        )
+
+        def authenticator_factory() -> Authenticator:
+            return get_async_proxy_authenticator(endpoint=endpoint, proxy_command=proxy_command, **kwargs)
+
+        return [
+            AuthUnaryUnaryInterceptor(authenticator_factory),
+            AuthUnaryStreamInterceptor(authenticator_factory),
+            AuthStreamUnaryInterceptor(authenticator_factory),
+            AuthStreamStreamInterceptor(authenticator_factory),
+        ]
+    else:
+        return []
+
+
+def get_async_authenticator(
+    endpoint: str,
+    cfg_store: ClientConfigStore,
+    *,
+    command: typing.Optional[typing.List[str]] = None,
+    insecure_skip_verify: bool = False,
+    auth_type: AuthType = "Pkce",
+    ca_cert_file_path: typing.Optional[str] = None,
+    **kwargs,
+) -> Authenticator:
+    """
+    Returns a new authenticator based on the platform config.
+    This is an async-compatible version of get_authenticator.
+    Must be async because it calls get_async_session which may perform IO operations.
+
+    :param endpoint: The endpoint URL for authentication
+    :param cfg_store: The client configuration store
+    :param command: Command to execute for ExternalCommand authentication
+    :param insecure_skip_verify: Whether to skip SSL certificate verification
+    :param auth_type: The authentication type to use
+    :param ca_cert_file_path: Path to CA certificate file for SSL verification
+    :param kwargs: Additional arguments passed to the authenticator, which may include:
+        - http_session: httpx.AsyncClient session to use for requests
+        - client_config: Optional client configuration containing authentication settings
+        - credentials: Optional credentials to use for authentication
+        - http_proxy_url: HTTP proxy URL
+        - verify: Whether to verify SSL certificates (bool or path to cert)
+        - ca_cert_path: Optional path to CA certificate file
+        - client_id: Client ID for ClientSecret authentication
+        - client_secret: Client secret for ClientSecret authentication (for ClientSecret auth)
+        - client_credentials_secret: Client secret for ClientSecret authentication (alias)
+        - scopes: List of scopes to request during authentication
+        - audience: Audience for the token
+        - header_key: Header key to use for authentication
+        - proxy_env: Environment variables for proxy command
+        - proxy_timeout: Timeout for proxy command execution
+        - redirect_uri: OAuth2 redirect URI for PKCE authentication
+        - add_request_auth_code_params_to_request_access_token_params: Whether to add auth code params to token request
+        - request_auth_code_params: Parameters to add to login URI opened in browser
+        - request_access_token_params: Parameters to add when exchanging auth code for access token
+        - refresh_access_token_params: Parameters to add when refreshing access token
+    :return: An authenticator instance
+    """
+    verify = None
+    if insecure_skip_verify:
+        verify = False
+    elif ca_cert_file_path:
+        verify = ca_cert_file_path
+
+    # Note: The authenticator classes already have async refresh_credentials methods
+    # so we can reuse them with our async session
+    match auth_type:
+        case "Pkce":
+            from union.remote._client.auth._authenticators.pkce import PKCEAuthenticator
+
+            return PKCEAuthenticator(endpoint=endpoint, cfg_store=cfg_store, verify=verify, **kwargs)
+        case "ClientSecret":
+            from union.remote._client.auth._authenticators.client_credentials import ClientCredentialsAuthenticator
+
+            return ClientCredentialsAuthenticator(endpoint=endpoint, cfg_store=cfg_store, verify=verify, **kwargs)
+        case "ExternalCommand":
+            from union.remote._client.auth._authenticators.external_command import AsyncCommandAuthenticator
+
+            return AsyncCommandAuthenticator(endpoint=endpoint, command=command, verify=verify, **kwargs)
+        case "DeviceFlow":
+            from union.remote._client.auth._authenticators.device_code import DeviceCodeAuthenticator
+
+            return DeviceCodeAuthenticator(endpoint=endpoint, cfg_store=cfg_store, verify=verify, **kwargs)
+        case _:
+            raise ValueError(
+                f"Invalid auth mode [{auth_type}] specified. Please update the creds config to use a valid value"
+            )
+
+
+def get_async_proxy_authenticator(endpoint: str, *, proxy_command: typing.List[str], **kwargs) -> Authenticator:
+    """
+    Returns an async authenticator for proxy authentication.
+    This function needs to be async because it calls get_async_command_authenticator which performs IO operations.
+
+    :param endpoint: The endpoint URL for authentication
+    :param proxy_command: Command to execute to get proxy authentication token
+    :param kwargs: Additional arguments passed to the authenticator, including:
+        - header_key: Header key to use for authentication (defaults to "proxy-authorization")
+        - proxy_env: Environment variables for the proxy command
+        - proxy_timeout: Timeout for the proxy command
+        - http_session: httpx.AsyncClient session to use for requests
+        - cfg_store: Optional client configuration store for retrieving remote configuration
+        - client_config: Optional client configuration containing authentication settings
+        - credentials: Optional credentials to use for authentication
+        - http_proxy_url: Optional HTTP proxy URL
+        - verify: Whether to verify SSL certificates (default: True)
+        - ca_cert_path: Optional path to CA certificate file
+    :return: An authenticator instance for proxy authentication
+    """
+    return AsyncCommandAuthenticator(
+        endpoint=endpoint, command=proxy_command, header_key="proxy-authorization", **kwargs
+    )