flyte 0.0.1b0__py3-none-any.whl

flyte/remote/_client/auth/_channel.py

@@ -0,0 +1,184 @@
+import ssl
+import typing
+
+import grpc
+import grpc.aio
+import httpx
+from grpc.experimental.aio import init_grpc_aio
+
+from flyte._logging import logger
+
+from ._authenticators.base import get_async_session
+from ._authenticators.factory import (
+    create_auth_interceptors,
+    create_proxy_auth_interceptors,
+    get_async_proxy_authenticator,
+)
+
+# Initialize gRPC AIO early enough so it can be used in the main thread
+init_grpc_aio()
+
+
+def bootstrap_ssl_from_server(endpoint: str) -> grpc.ChannelCredentials:
+    """
+    Retrieves the SSL certificate from the remote server and creates gRPC channel credentials.
+
+    This function should be used only when insecure-skip-verify is enabled. It extracts the server address
+    and port from the endpoint URL, retrieves the SSL certificate from the server, and creates
+    gRPC channel credentials using the certificate.
+
+    :param endpoint: The endpoint URL to retrieve the SSL certificate from, may include port number
+    :return: gRPC channel credentials created from the retrieved certificate
+    """
+    # Get port from endpoint or use 443
+    endpoint_parts = endpoint.rsplit(":", 1)
+    if len(endpoint_parts) == 2 and endpoint_parts[1].isdigit():
+        server_address = (endpoint_parts[0], int(endpoint_parts[1]))
+    else:
+        logger.warning(f"Unrecognized port in endpoint [{endpoint}], defaulting to 443.")
+        server_address = (endpoint, 443)
+
+    # Run the blocking SSL certificate retrieval in a thread pool
+    cert = ssl.get_server_certificate(server_address)
+    return grpc.ssl_channel_credentials(str.encode(cert))
+
+
+async def create_channel(
+    endpoint: str,
+    insecure: typing.Optional[bool] = None,
+    insecure_skip_verify: typing.Optional[bool] = False,
+    ca_cert_file_path: typing.Optional[str] = None,
+    ssl_credentials: typing.Optional[grpc.ssl_channel_credentials] = None,
+    grpc_options: typing.Optional[typing.Sequence[typing.Tuple[str, typing.Any]]] = None,
+    compression: typing.Optional[grpc.Compression] = None,
+    http_session: httpx.AsyncClient | None = None,
+    proxy_command: typing.List[str] | None = None,
+    **kwargs,
+) -> grpc.aio.Channel:
+    """
+    Creates a new gRPC channel with appropriate authentication interceptors.
+
+    This function creates either a secure or insecure gRPC channel based on the provided parameters,
+    and adds authentication interceptors to the channel. If SSL credentials are not provided,
+    they are created based on the insecure_skip_verify and ca_cert_file_path parameters.
+
+    The function is async because it may need to read certificate files asynchronously
+    and create authentication interceptors that perform async operations.
+
+    :param endpoint: The endpoint URL for the gRPC channel
+    :param insecure: Whether to use an insecure channel (no SSL)
+    :param insecure_skip_verify: Whether to skip SSL certificate verification
+    :param ca_cert_file_path: Path to CA certificate file for SSL verification
+    :param ssl_credentials: Pre-configured SSL credentials for the channel
+    :param grpc_options: Additional gRPC channel options
+    :param compression: Compression method for the channel
+    :param http_session: Pre-configured HTTP session to use for requests
+    :param proxy_command: List of strings for proxy command configuration
+    :param kwargs: Additional arguments passed to various functions:
+        - For grpc.aio.insecure_channel/secure_channel:
+            - root_certificates: Root certificates for SSL credentials
+            - private_key: Private key for SSL credentials
+            - certificate_chain: Certificate chain for SSL credentials
+            - options: gRPC channel options
+            - compression: gRPC compression method
+        - For proxy configuration:
+            - proxy_env: Dict of environment variables for proxy
+            - proxy_timeout: Timeout for proxy connection
+        - For authentication interceptors (passed to create_auth_interceptors and create_proxy_auth_interceptors):
+            - 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
+            - client_credentials_secret: Client secret for ClientSecret authentication (alias)
+            - scopes: List of scopes to request during authentication
+            - audience: Audience for the token
+            - http_proxy_url: HTTP proxy URL
+            - verify: Whether to verify SSL certificates
+            - ca_cert_path: Optional path to CA certificate file
+            - header_key: Header key to use for authentication
+            - 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: grpc.aio.Channel with authentication interceptors configured
+    """
+
+    if not ssl_credentials:
+        if insecure_skip_verify:
+            ssl_credentials = bootstrap_ssl_from_server(endpoint)
+        elif ca_cert_file_path:
+            import aiofiles
+
+            async with aiofiles.open(ca_cert_file_path, "rb") as f:
+                st_cert = f.read()
+            ssl_credentials = grpc.ssl_channel_credentials(st_cert)
+        else:
+            ssl_credentials = grpc.ssl_channel_credentials()
+
+    # Create an unauthenticated channel first to use to get the server metadata
+    if insecure:
+        unauthenticated_channel = grpc.aio.insecure_channel(endpoint, **kwargs)
+    else:
+        unauthenticated_channel = grpc.aio.secure_channel(
+            target=endpoint,
+            credentials=ssl_credentials,
+            options=grpc_options,
+            compression=compression,
+        )
+
+    from ._grpc_utils.default_metadata_interceptor import (
+        DefaultMetadataStreamStreamInterceptor,
+        DefaultMetadataStreamUnaryInterceptor,
+        DefaultMetadataUnaryStreamInterceptor,
+        DefaultMetadataUnaryUnaryInterceptor,
+    )
+
+    # Add all types of default metadata interceptors
+    interceptors: typing.List[grpc.aio.ClientInterceptor] = [
+        DefaultMetadataUnaryUnaryInterceptor(),
+        DefaultMetadataUnaryStreamInterceptor(),
+        DefaultMetadataStreamUnaryInterceptor(),
+        DefaultMetadataStreamStreamInterceptor(),
+    ]
+
+    # Create an HTTP session if not provided so we share the same http client across the stack
+    if not http_session:
+        proxy_authenticator = None
+        if proxy_command:
+            proxy_authenticator = get_async_proxy_authenticator(
+                endpoint=endpoint, proxy_command=proxy_command, **kwargs
+            )
+
+        http_session = get_async_session(
+            ca_cert_file_path=ca_cert_file_path, proxy_authenticator=proxy_authenticator, **kwargs
+        )
+
+    # Get proxy auth interceptors
+    proxy_auth_interceptors = create_proxy_auth_interceptors(endpoint, http_session=http_session, **kwargs)
+    interceptors.extend(proxy_auth_interceptors)
+
+    # Get auth interceptors
+    auth_interceptors = create_auth_interceptors(
+        endpoint=endpoint,
+        in_channel=unauthenticated_channel,
+        insecure=insecure,
+        insecure_skip_verify=insecure_skip_verify,
+        ca_cert_file_path=ca_cert_file_path,
+        http_session=http_session,
+        **kwargs,
+    )
+
+    interceptors.extend(auth_interceptors)
+
+    if insecure:
+        return grpc.aio.insecure_channel(endpoint, interceptors=interceptors, **kwargs)
+
+    return grpc.aio.secure_channel(
+        target=endpoint,
+        credentials=ssl_credentials,
+        options=grpc_options,
+        compression=compression,
+        interceptors=interceptors,
+    )

flyte/remote/_client/auth/_client_config.py

@@ -0,0 +1,83 @@
+import typing
+from abc import abstractmethod
+
+import grpc.aio
+import pydantic
+from flyteidl.service.auth_pb2 import OAuth2MetadataRequest, PublicClientAuthConfigRequest
+from flyteidl.service.auth_pb2_grpc import AuthMetadataServiceStub
+
+AuthType = typing.Literal["ClientSecret", "Pkce", "ExternalCommand", "DeviceFlow"]
+
+
+class ClientConfig(pydantic.BaseModel):
+    """
+    Client Configuration that is needed by the authenticator
+    """
+
+    token_endpoint: str
+    authorization_endpoint: str
+    redirect_uri: str
+    client_id: str
+    device_authorization_endpoint: typing.Optional[str] = None
+    scopes: typing.Optional[typing.List[str]] = None
+    header_key: str = "authorization"
+    audience: typing.Optional[str] = None
+
+    def with_override(self, other: "ClientConfig") -> "ClientConfig":
+        """
+        Returns a new ClientConfig instance with the values from the other instance overriding the current instance.
+        """
+        return ClientConfig(
+            token_endpoint=other.token_endpoint or self.token_endpoint,
+            authorization_endpoint=other.authorization_endpoint or self.authorization_endpoint,
+            redirect_uri=other.redirect_uri or self.redirect_uri,
+            client_id=other.client_id or self.client_id,
+            device_authorization_endpoint=other.device_authorization_endpoint or self.device_authorization_endpoint,
+            scopes=other.scopes or self.scopes,
+            header_key=other.header_key or self.header_key,
+            audience=other.audience or self.audience,
+        )
+
+
+class ClientConfigStore(object):
+    """
+    Client Config store retrieve client config. this can be done in multiple ways
+    """
+
+    @abstractmethod
+    async def get_client_config(self) -> ClientConfig: ...
+
+
+class StaticClientConfigStore(ClientConfigStore):
+    def __init__(self, cfg: ClientConfig):
+        self._cfg = cfg
+
+    async def get_client_config(self) -> ClientConfig:
+        return self._cfg
+
+
+class RemoteClientConfigStore(ClientConfigStore):
+    """
+    This class implements the ClientConfigStore that is served by the Flyte Server, that implements AuthMetadataService
+    """
+
+    def __init__(self, unauthenticated_channel: grpc.aio.Channel):
+        self._unauthenticated_channel = unauthenticated_channel
+
+    async def get_client_config(self) -> ClientConfig:
+        """
+        Retrieves the ClientConfig from the given grpc.Channel assuming  AuthMetadataService is available
+        """
+        metadata_service = AuthMetadataServiceStub(self._unauthenticated_channel)
+        public_client_config = await metadata_service.GetPublicClientConfig(PublicClientAuthConfigRequest())
+        oauth2_metadata = await metadata_service.GetOAuth2Metadata(OAuth2MetadataRequest())
+        return ClientConfig(
+            token_endpoint=oauth2_metadata.token_endpoint,
+            authorization_endpoint=oauth2_metadata.authorization_endpoint,
+            redirect_uri=public_client_config.redirect_uri,
+            client_id=public_client_config.client_id,
+            scopes=public_client_config.scopes,
+            header_key=public_client_config.authorization_metadata_key,
+            device_authorization_endpoint=oauth2_metadata.device_authorization_endpoint,
+            audience=public_client_config.audience,
+        )

flyte/remote/_client/auth/_default_html.py

@@ -0,0 +1,32 @@
+import textwrap
+
+
+def get_default_success_html(endpoint: str) -> str:
+    """Get default success html."""
+    return textwrap.dedent(
+        """
+    <html>
+    <head>
+        <title>OAuth2 Authentication to Union Successful</title>
+    </head>
+    <body style="background:white;font-family:Arial">
+        <div style="position: absolute;top:40%;left:50%;transform: translate(-50%, -50%);text-align:center;">
+            <div style="margin:auto">
+                <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 300 65" fill="currentColor"
+                    style="color:#fdb51e;width:360px;">
+                    <title>Union.ai</title>
+                    <path d="M32,64.8C14.4,64.8,0,51.5,0,34V3.6h17.6v41.3c0,1.9,1.1,3,3,3h23c1.9,0,3-1.1,3-3V3.6H64V34
+                    C64,51.5,49.6,64.8,32,64.8z M69.9,30.9v30.4h17.6V20c0-1.9,1.1-3,3-3h23c1.9,0,3,1.1,3,3v41.3H134V30.9c0-17.5-14.4-30.8-32.1-30.8
+                    S69.9,13.5,69.9,30.9z M236,30.9v30.4h17.6V20c0-1.9,1.1-3,3-3h23c1.9,0,3,1.1,3,3v41.3H300V30.9c0-17.5-14.4-30.8-32-30.8
+                    S236,13.5,236,30.9L236,30.9z M230.1,32.4c0,18.2-14.2,32.5-32.2,32.5s-32-14.3-32-32.5s14-32.1,32-32.1S230.1,14.3,230.1,32.4
+                    L230.1,32.4z M213.5,20.2c0-1.9-1.1-3-3-3h-24.8c-1.9,0-3,1.1-3,3v24.5c0,1.9,1.1,3,3,3h24.8c1.9,0,3-1.1,3-3V20.2z M158.9,3.6
+                    h-17.6v57.8h17.6V3.6z"></path>
+                </svg>
+                <h2>You've successfully authenticated to Union!</h2>
+                <p style="font-size:20px;">Return to your terminal for next steps</p>
+            </div>
+        </div>
+    </body>
+    </html>
+    """  # noqa: E501
+    )

flyte/remote/_client/auth/_grpc_utils/auth_interceptor.py

@@ -0,0 +1,288 @@
+import typing
+from typing import AsyncIterator, Optional, Union
+
+import grpc.aio
+from grpc.aio import ClientCallDetails, Metadata
+from grpc.aio._typing import DoneCallbackType, EOFType, RequestType, ResponseType
+
+from flyte.remote._client.auth._authenticators.base import Authenticator
+from flyte.remote._client.auth._grpc_utils.default_metadata_interceptor import with_metadata
+
+
+class _BaseAuthInterceptor:
+    """
+    Base class for all auth interceptors that provides common authentication functionality.
+    """
+
+    def __init__(self, get_authenticator: typing.Callable[[], Authenticator]):
+        self._get_authenticator = get_authenticator
+        self._authenticator: typing.Optional[Authenticator] = None
+
+    @property
+    def authenticator(self) -> Authenticator:
+        if self._authenticator is None:
+            self._authenticator = self._get_authenticator()
+        return self._authenticator
+
+    async def call_details_with_auth_metadata(
+        self, client_call_details: grpc.aio.ClientCallDetails
+    ) -> typing.Tuple[grpc.aio.ClientCallDetails, str]:
+        """
+        Returns new ClientCallDetails with authentication metadata added.
+
+        This method retrieves authentication metadata from the authenticator and adds it to the
+        client call details. If no authentication metadata is available, the original client call
+        details are returned unchanged.
+
+        :param client_call_details: The original client call details containing method, timeout, metadata,
+                                   credentials, and wait_for_ready settings
+        :return: Updated client call details with authentication metadata added to the existing metadata
+        """
+        auth_metadata = await self.authenticator.get_grpc_call_auth_metadata()
+        if auth_metadata:
+            return with_metadata(client_call_details, auth_metadata.pairs), auth_metadata.creds_id
+        else:
+            return client_call_details, ""
+
+
+class AuthUnaryUnaryInterceptor(_BaseAuthInterceptor, grpc.aio.UnaryUnaryClientInterceptor):
+    """
+    Interceptor for unary-unary RPC calls that adds authentication metadata.
+    """
+
+    async def intercept_unary_unary(
+        self,
+        continuation: typing.Callable,
+        client_call_details: ClientCallDetails,
+        request: typing.Any,
+    ):
+        """
+        Intercepts unary-unary calls and adds auth metadata if available. On Unauthenticated, resets the token and
+        refreshes and then retries with the new token.
+
+        This method first adds authentication metadata to the client call details, then attempts to make the RPC call.
+        If the call fails with an UNAUTHENTICATED or UNKNOWN status code, it refreshes the credentials and retries
+        the call with the new authentication metadata.
+
+        :param continuation: Function to continue the RPC call chain with the updated call details
+        :param client_call_details: Details about the RPC call including method, timeout, metadata, credentials,
+        and wait_for_ready
+        :param request: The request message to be sent to the server
+        :return: The response from the RPC call after successful authentication
+        :raises: grpc.aio.AioRpcError if the call fails for reasons other than authentication
+        """
+        updated_call_details, creds_id = await self.call_details_with_auth_metadata(client_call_details)
+        try:
+            return await (await continuation(updated_call_details, request))
+        except grpc.aio.AioRpcError as e:
+            if e.code() == grpc.StatusCode.UNAUTHENTICATED or e.code() == grpc.StatusCode.UNKNOWN:
+                await self.authenticator.refresh_credentials(creds_id=creds_id)
+                updated_call_details, _ = await self.call_details_with_auth_metadata(client_call_details)
+                return await (await continuation(updated_call_details, request))
+            else:
+                raise e
+
+
+class UnaryStreamCall(grpc.aio.UnaryStreamCall):
+    def __init__(
+        self,
+        parent_interceptor: _BaseAuthInterceptor,
+        authenticator: Authenticator,
+        continuation: typing.Callable,
+        call_details: grpc.aio.ClientCallDetails,
+        request: RequestType,
+    ):
+        super().__init__()
+        self._continuation = continuation
+        self._call_details = call_details
+        self._request = request
+        self._authenticator = authenticator
+        self._parent_interceptor = parent_interceptor
+        self._call: (
+            Union[
+                grpc.aio.UnaryStreamCall[RequestType, ResponseType],
+                grpc.aio.StreamStreamCall[RequestType, ResponseType],
+            ]
+            | None
+        ) = None
+
+    async def response_iterator(self) -> typing.AsyncIterator[ResponseType]:
+        call_details, creds_id = await self._parent_interceptor.call_details_with_auth_metadata(self._call_details)
+        self._call = await self._continuation(call_details, self._request)
+        try:
+            async for response in self._call:
+                yield response
+        except grpc.aio.AioRpcError as e:
+            if e.code() == grpc.StatusCode.UNAUTHENTICATED or e.code() == grpc.StatusCode.UNKNOWN:
+                await self._authenticator.refresh_credentials(creds_id=creds_id)
+                updated_call_details, _ = await self._parent_interceptor.call_details_with_auth_metadata(call_details)
+                self._call = await self._continuation(updated_call_details, self._request)
+                async for response in self._call:
+                    yield response
+            else:
+                raise e
+
+    def __aiter__(self) -> AsyncIterator[ResponseType]:
+        return self.response_iterator()
+
+    async def read(self) -> Union[EOFType, ResponseType]:
+        if self._call is not None:
+            return await self._call.read()
+        return EOFType()
+
+    async def initial_metadata(self) -> Metadata:
+        if self._call is not None:
+            return await self._call.initial_metadata()
+        return Metadata()
+
+    async def trailing_metadata(self) -> Metadata:
+        if self._call is not None:
+            return await self._call.trailing_metadata()
+        return Metadata()
+
+    async def code(self) -> grpc.StatusCode:
+        if self._call is not None:
+            return await self._call.code()
+        return grpc.StatusCode.OK
+
+    async def details(self) -> str:
+        if self._call is not None:
+            return await self._call.details()
+        return ""
+
+    async def wait_for_connection(self) -> None:
+        if self._call is not None:
+            await self._call.wait_for_connection()
+        return None
+
+    def cancelled(self) -> bool:
+        if self._call is not None:
+            return self._call.cancelled()
+        return False
+
+    def done(self) -> bool:
+        if self._call is not None:
+            return self._call.done()
+        return False
+
+    def time_remaining(self) -> Optional[float]:
+        if self._call is not None:
+            return self._call.time_remaining()
+        return None
+
+    def cancel(self) -> bool:
+        if self._call is not None:
+            return self._call.cancel()
+        return False
+
+    def add_done_callback(self, callback: DoneCallbackType) -> None:
+        if self._call is not None:
+            self._call.add_done_callback(callback=callback)
+        return None
+
+
+class AuthUnaryStreamInterceptor(_BaseAuthInterceptor, grpc.aio.UnaryStreamClientInterceptor):
+    """
+    Interceptor for unary-stream RPC calls that adds authentication metadata.
+    """
+
+    async def intercept_unary_stream(
+        self, continuation: typing.Callable, client_call_details: grpc.aio.ClientCallDetails, request: typing.Any
+    ):
+        """
+        Intercepts unary-stream calls and adds auth metadata if available.
+
+        This method first adds authentication metadata to the client call details, then attempts to make the RPC call.
+        If the call fails with an UNAUTHENTICATED or UNKNOWN status code, it refreshes the credentials and retries
+        the call with the new authentication metadata.
+
+        :param continuation: Function to continue the RPC call chain with the updated call details
+        :param client_call_details: Details about the RPC call including method, timeout, metadata, credentials,
+        and wait_for_ready
+        :param request: The request message to be sent to the server
+        :return: A stream of responses from the RPC call after successful authentication
+        :raises: grpc.aio.AioRpcError if the call fails for reasons other than authentication
+        """
+
+        return UnaryStreamCall(
+            parent_interceptor=self,
+            authenticator=self.authenticator,
+            call_details=client_call_details,
+            continuation=continuation,
+            request=request,
+        )
+
+
+class AuthStreamUnaryInterceptor(_BaseAuthInterceptor, grpc.aio.StreamUnaryClientInterceptor):
+    """
+    Interceptor for stream-unary RPC calls that adds authentication metadata.
+    """
+
+    async def intercept_stream_unary(
+        self,
+        continuation: typing.Callable,
+        client_call_details: grpc.aio.ClientCallDetails,
+        request_iterator: typing.Any,
+    ):
+        """
+        Intercepts stream-unary calls and adds auth metadata if available.
+
+        This method first adds authentication metadata to the client call details, then attempts to make the RPC call.
+        If the call fails with an UNAUTHENTICATED or UNKNOWN status code, it refreshes the credentials and retries
+        the call with the new authentication metadata.
+
+        :param continuation: Function to continue the RPC call chain with the updated call details
+        :param client_call_details: Details about the RPC call including method, timeout, metadata, credentials,
+        and wait_for_ready
+        :param request_iterator: An iterator of request messages to be sent to the server
+        :return: The response from the RPC call after successful authentication
+        :raises: grpc.aio.AioRpcError if the call fails for reasons other than authentication
+        """
+        updated_call_details, creds_id = await self.call_details_with_auth_metadata(client_call_details)
+        try:
+            call = await continuation(updated_call_details, request_iterator)
+            return await call
+        except grpc.aio.AioRpcError as e:
+            if e.code() == grpc.StatusCode.UNAUTHENTICATED or e.code() == grpc.StatusCode.UNKNOWN:
+                await self.authenticator.refresh_credentials(creds_id=creds_id)
+                updated_call_details, _ = await self.call_details_with_auth_metadata(client_call_details)
+                call = await continuation(updated_call_details, request_iterator)
+                return await call
+            raise e
+
+
+class AuthStreamStreamInterceptor(_BaseAuthInterceptor, grpc.aio.StreamStreamClientInterceptor):
+    """
+    Interceptor for stream-stream RPC calls that adds authentication metadata.
+    """
+
+    async def intercept_stream_stream(
+        self,
+        continuation: typing.Callable,
+        client_call_details: grpc.aio.ClientCallDetails,
+        request_iterator: typing.Any,
+    ):
+        """
+        Intercepts stream-stream calls and adds auth metadata if available.
+
+        This method first adds authentication metadata to the client call details, then attempts to make the RPC call.
+        If the call fails with an UNAUTHENTICATED or UNKNOWN status code, it refreshes the credentials and retries
+        the call with the new authentication metadata.
+
+        :param continuation: Function to continue the RPC call chain with the updated call details
+        :param client_call_details: Details about the RPC call including method, timeout, metadata, credentials,
+        and wait_for_ready
+        :param request_iterator: An iterator of request messages to be sent to the server
+        :return: A stream of responses from the RPC call after successful authentication
+        """
+        return UnaryStreamCall(
+            parent_interceptor=self,
+            authenticator=self.authenticator,
+            call_details=client_call_details,
+            continuation=continuation,
+            request=request_iterator,
+        )
+
+
+# For backward compatibility, maintain the original class name but as a type alias
+AuthUnaryInterceptor = AuthUnaryUnaryInterceptor