diff --git a/flytekit/clients/auth/keyring.py b/flytekit/clients/auth/keyring.py index 2aa4dc0d687..2d4b4488f06 100644 --- a/flytekit/clients/auth/keyring.py +++ b/flytekit/clients/auth/keyring.py @@ -62,7 +62,7 @@ def retrieve(for_endpoint: str) -> typing.Optional[Credentials]: logging.debug(f"KeyRing not available, tokens will not be cached. Error: {e}") return None - if not access_token: + if not access_token and not id_token: return None return Credentials(access_token, refresh_token, for_endpoint, id_token=id_token) diff --git a/plugins/flytekit-identity-aware-proxy/flytekitplugins/identity_aware_proxy/cli.py b/plugins/flytekit-identity-aware-proxy/flytekitplugins/identity_aware_proxy/cli.py index cf98169799c..65f7f340e50 100644 --- a/plugins/flytekit-identity-aware-proxy/flytekitplugins/identity_aware_proxy/cli.py +++ b/plugins/flytekit-identity-aware-proxy/flytekitplugins/identity_aware_proxy/cli.py @@ -1,14 +1,26 @@ import logging +import os import typing import click +import jwt from google.api_core.exceptions import NotFound +from google.auth import default +from google.auth.transport.requests import Request from google.cloud import secretmanager +from google.oauth2 import id_token from flytekit.clients.auth.auth_client import AuthorizationClient from flytekit.clients.auth.authenticator import Authenticator from flytekit.clients.auth.exceptions import AccessTokenNotFoundError -from flytekit.clients.auth.keyring import KeyringStore +from flytekit.clients.auth.keyring import Credentials, KeyringStore + +WEBAPP_CLIENT_ID_HELP = ( + "Webapp type OAuth 2.0 client ID used by the IAP. " + "Typically in the form of `.apps.googleusercontent.com`. " + "Created when activating IAP for the Flyte deployment. " + "https://cloud.google.com/iap/docs/enabling-kubernetes-howto#oauth-credentials" +) class GCPIdentityAwareProxyAuthenticator(Authenticator): @@ -131,11 +143,7 @@ def cli(): type=str, default=None, required=True, - help=( - "Webapp type OAuth 2.0 client ID. Typically in the form of `.apps.googleusercontent.com`. " - "Created when activating IAP for the Flyte deployment. " - "https://cloud.google.com/iap/docs/enabling-kubernetes-howto#oauth-credentials" - ), + help=WEBAPP_CLIENT_ID_HELP, ) @click.option( "--project", @@ -147,7 +155,7 @@ def cli(): def generate_user_id_token( desktop_client_id: str, desktop_client_secret_gcp_secret_name: str, webapp_client_id: str, project: str ): - """Generate a user account ID token for proxy-authentication/authorization with GCP Identity Aware Proxy.""" + """Generate a user account ID token for proxy-authorization with GCP Identity Aware Proxy.""" desktop_client_secret = get_gcp_secret_manager_secret(project, desktop_client_secret_gcp_secret_name) iap_authenticator = GCPIdentityAwareProxyAuthenticator( @@ -163,5 +171,76 @@ def generate_user_id_token( click.echo(iap_authenticator.get_credentials().id_token) +def get_service_account_id_token(audience: str) -> str: + """Fetch an ID Token for the service account used by the current environment. + + Uses flytekit's KeyringStore to cache the ID token. + + This function acquires ID token from the environment in the following order. + See https://google.aip.dev/auth/4110. + + 1. If the environment variable ``GOOGLE_APPLICATION_CREDENTIALS`` is set + to the path of a valid service account JSON file, then ID token is + acquired using this service account credentials. + 2. If the application is running in Compute Engine, App Engine or Cloud Run, + then the ID token are obtained from the metadata server. + + Args: + audience (str): The audience that this ID token is intended for. + """ + credentials, _ = default() + # Flytekit's KeyringStore, by default, uses the endpoint as the key to store the credentials + # We use the audience and the service account email as the key + audience_and_account_key = audience + "-" + credentials.service_account_email + creds = KeyringStore.retrieve(audience_and_account_key) + if creds: + is_expired = False + try: + exp_margin = -300 # Generate a new token if it expires in less than 5 minutes + jwt.decode( + creds.id_token.encode("utf-8"), + options={"verify_signature": False, "verify_exp": True}, + leeway=exp_margin, + ) + except jwt.ExpiredSignatureError: + is_expired = True + + if not is_expired: + return creds.id_token + + token = id_token.fetch_id_token(Request(), audience) + + KeyringStore.store(Credentials(for_endpoint=audience_and_account_key, access_token="", id_token=token)) + return token + + +@cli.command() +@click.option( + "--webapp_client_id", + type=str, + default=None, + required=True, + help=WEBAPP_CLIENT_ID_HELP, +) +@click.option( + "--service_account_key", + type=click.Path(exists=True, dir_okay=False), + default=None, + required=False, + help=( + "Path to a service account key file. Alternatively set the environment variable " + "`GOOGLE_APPLICATION_CREDENTIALS` to the path of the service account key file. " + "If not provided and in Compute Engine, App Engine, or Cloud Run, will retrieve " + "the ID token from the metadata server." + ), +) +def generate_service_account_id_token(webapp_client_id: str, service_account_key: str): + """Generate a service account ID token for proxy-authorization with GCP Identity Aware Proxy.""" + if service_account_key: + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = service_account_key + token = get_service_account_id_token(webapp_client_id) + click.echo(token) + + if __name__ == "__main__": cli()