Feat: Add plugin for generating GCP IAP ID tokens via external command

flytekit/clients/auth/auth_client.py

@@ -185,6 +185,10 @@ def __init__(
         endpoint_metadata: typing.Optional[EndpointMetadata] = None,
         verify: typing.Optional[typing.Union[bool, str]] = None,
         session: typing.Optional[_requests.Session] = None,
+        request_auth_code_params: typing.Optional[typing.Dict[str, str]] = None,
+        request_access_token_params: typing.Optional[typing.Dict[str, str]] = None,
+        refresh_access_token_params: typing.Optional[typing.Dict[str, str]] = None,
+        add_request_auth_code_params_to_request_access_token_params: typing.Optional[bool] = False,
     ):
         """
         Create new AuthorizationClient
@@ -193,7 +197,9 @@ def __init__(
         :param auth_endpoint: str endpoint where auth metadata can be found
         :param token_endpoint: str endpoint to retrieve token from
         :param scopes: list[str] oauth2 scopes
-        :param client_id
+        :param client_id: oauth2 client id
+        :param redirect_uri: oauth2 redirect uri
+        :param endpoint_metadata: EndpointMetadata object to control the rendering of the page on login successful or failure
         :param verify: (optional) Either a boolean, in which case it controls whether we verify
             the server's TLS certificate, or a string, in which case it must be a path
             to a CA bundle to use. Defaults to ``True``. When set to
@@ -204,6 +210,13 @@ def __init__(
             may be useful during local development or testing.
         :param session: (optional) A custom requests.Session object to use for making HTTP requests.
             If not provided, a new Session object will be created.
+        :param request_auth_code_params: (optional) dict of parameters to add to login uri opened in the browser
+        :param request_access_token_params: (optional) dict of parameters to add when exchanging the auth code for the access token
+        :param refresh_access_token_params: (optional) dict of parameters to add when refreshing the access token
+        :param add_request_auth_code_params_to_request_access_token_params: Whether to add the `request_auth_code_params` to
+            the parameters sent when exchanging the auth code for the access token. Defaults to False.
+            Required e.g. for the PKCE flow with flyteadmin.
+            Not required for e.g. the standard OAuth2 flow on GCP.
         """
         self._endpoint = endpoint
         self._auth_endpoint = auth_endpoint
@@ -216,16 +229,13 @@ def __init__(
         self._client_id = client_id
         self._scopes = scopes or []
         self._redirect_uri = redirect_uri
-        self._code_verifier = _generate_code_verifier()
-        code_challenge = _create_code_challenge(self._code_verifier)
-        self._code_challenge = code_challenge
         state = _generate_state_parameter()
         self._state = state
         self._verify = verify
         self._headers = {"content-type": "application/x-www-form-urlencoded"}
         self._session = session or _requests.Session()
 
-        self._params = {
+        self._request_auth_code_params = {
             "client_id": client_id,  # This must match the Client ID of the OAuth application.
             "response_type": "code",  # Indicates the authorization code grant
             "scope": " ".join(s.strip("' ") for s in self._scopes).strip(
@@ -234,10 +244,18 @@ def __init__(
             # callback location where the user-agent will be directed to.
             "redirect_uri": self._redirect_uri,
             "state": state,
-            "code_challenge": code_challenge,
-            "code_challenge_method": "S256",
         }
 
+        if request_auth_code_params:
+            # Allow adding additional parameters to the request_auth_code_params
+            self._request_auth_code_params.update(request_auth_code_params)
+
+        self._request_access_token_params = request_access_token_params or {}
+        self._refresh_access_token_params = refresh_access_token_params or {}
+
+        if add_request_auth_code_params_to_request_access_token_params:
+            self._request_access_token_params.update(self._request_auth_code_params)
+
     def __repr__(self):
         return f"AuthorizationClient({self._auth_endpoint}, {self._token_endpoint}, {self._client_id}, {self._scopes}, {self._redirect_uri})"
 
@@ -253,7 +271,7 @@ def _create_callback_server(self):
 
     def _request_authorization_code(self):
         scheme, netloc, path, _, _, _ = _urlparse.urlparse(self._auth_endpoint)
-        query = _urlencode(self._params)
+        query = _urlencode(self._request_auth_code_params)
         endpoint = _urlparse.urlunparse((scheme, netloc, path, None, query, None))
         logging.debug(f"Requesting authorization code through {endpoint}")
         _webbrowser.open_new_tab(endpoint)
@@ -266,33 +284,38 @@ def _credentials_from_response(self, auth_token_resp) -> Credentials:
           "refresh_token": "bar",
           "token_type": "Bearer"
         }
+
+        Can additionally contain "expires_in" and "id_token" fields.
         """
         response_body = auth_token_resp.json()
         refresh_token = None
+        id_token = None
         if "access_token" not in response_body:
             raise ValueError('Expected "access_token" in response from oauth server')
         if "refresh_token" in response_body:
             refresh_token = response_body["refresh_token"]
         if "expires_in" in response_body:
             expires_in = response_body["expires_in"]
         access_token = response_body["access_token"]
+        if "id_token" in response_body:
+            id_token = response_body["id_token"]
 
-        return Credentials(access_token, refresh_token, self._endpoint, expires_in=expires_in)
+        return Credentials(access_token, refresh_token, self._endpoint, expires_in=expires_in, id_token=id_token)
 
     def _request_access_token(self, auth_code) -> Credentials:
         if self._state != auth_code.state:
             raise ValueError(f"Unexpected state parameter [{auth_code.state}] passed")
-        self._params.update(
-            {
-                "code": auth_code.code,
-                "code_verifier": self._code_verifier,
-                "grant_type": "authorization_code",
-            }
-        )
+
+        params = {
+            "code": auth_code.code,
+            "grant_type": "authorization_code",
+        }
+
+        params.update(self._request_access_token_params)
 
         resp = self._session.post(
             url=self._token_endpoint,
-            data=self._params,
+            data=params,
             headers=self._headers,
             allow_redirects=False,
             verify=self._verify,
@@ -336,13 +359,17 @@ def refresh_access_token(self, credentials: Credentials) -> Credentials:
         if credentials.refresh_token is None:
             raise ValueError("no refresh token available with which to refresh authorization credentials")
 
+        data = {
+            "refresh_token": credentials.refresh_token,
+            "grant_type": "refresh_token",
+            "client_id": self._client_id,
+        }
+
+        data.update(self._refresh_access_token_params)
+
         resp = self._session.post(
             url=self._token_endpoint,
-            data={
-                "grant_type": "refresh_token",
-                "client_id": self._client_id,
-                "refresh_token": credentials.refresh_token,
-            },
+            data=data,
             headers=self._headers,
             allow_redirects=False,
             verify=self._verify,

flytekit/clients/auth/authenticator.py

@@ -108,6 +108,12 @@ def __init__(
 
     def _initialize_auth_client(self):
         if not self._auth_client:
+
+            from .auth_client import _create_code_challenge, _generate_code_verifier
+
+            code_verifier = _generate_code_verifier()
+            code_challenge = _create_code_challenge(code_verifier)
+
             cfg = self._cfg_store.get_client_config()
             self._set_header_key(cfg.header_key)
             self._auth_client = AuthorizationClient(
@@ -119,6 +125,15 @@ def _initialize_auth_client(self):
                 token_endpoint=cfg.token_endpoint,
                 verify=self._verify,
                 session=self._session,
+                request_auth_code_params={
+                    "code_challenge": code_challenge,
+                    "code_challenge_method": "S256",
+                },
+                request_access_token_params={
+                    "code_verifier": code_verifier,
+                },
+                refresh_access_token_params={},
+                add_request_auth_code_params_to_request_access_token_params=True,
             )
 
     def refresh_credentials(self):

flytekit/clients/auth/keyring.py

@@ -3,7 +3,7 @@
 from dataclasses import dataclass
 
 import keyring as _keyring
-from keyring.errors import NoKeyringError
+from keyring.errors import NoKeyringError, PasswordDeleteError
 
 
 @dataclass
@@ -16,6 +16,7 @@ class Credentials(object):
     refresh_token: str = "na"
     for_endpoint: str = "flyte-default"
     expires_in: typing.Optional[int] = None
+    id_token: typing.Optional[str] = None
 
 
 class KeyringStore:
@@ -25,20 +26,28 @@ class KeyringStore:
 
     _access_token_key = "access_token"
     _refresh_token_key = "refresh_token"
+    _id_token_key = "id_token"
 
     @staticmethod
     def store(credentials: Credentials) -> Credentials:
         try:
-            _keyring.set_password(
-                credentials.for_endpoint,
-                KeyringStore._refresh_token_key,
-                credentials.refresh_token,
-            )
+            if credentials.refresh_token:
+                _keyring.set_password(
+                    credentials.for_endpoint,
+                    KeyringStore._refresh_token_key,
+                    credentials.refresh_token,
+                )
             _keyring.set_password(
                 credentials.for_endpoint,
                 KeyringStore._access_token_key,
                 credentials.access_token,
             )
+            if credentials.id_token:
+                _keyring.set_password(
+                    credentials.for_endpoint,
+                    KeyringStore._id_token_key,
+                    credentials.id_token,
+                )
         except NoKeyringError as e:
             logging.debug(f"KeyRing not available, tokens will not be cached. Error: {e}")
         return credentials
@@ -48,18 +57,23 @@ def retrieve(for_endpoint: str) -> typing.Optional[Credentials]:
         try:
             refresh_token = _keyring.get_password(for_endpoint, KeyringStore._refresh_token_key)
             access_token = _keyring.get_password(for_endpoint, KeyringStore._access_token_key)
+            id_token = _keyring.get_password(for_endpoint, KeyringStore._id_token_key)
         except NoKeyringError as e:
             logging.debug(f"KeyRing not available, tokens will not be cached. Error: {e}")
             return None
 
         if not access_token:
             return None
-        return Credentials(access_token, refresh_token, for_endpoint)
+        return Credentials(access_token, refresh_token, for_endpoint, id_token=id_token)
 
     @staticmethod
     def delete(for_endpoint: str):
         try:
             _keyring.delete_password(for_endpoint, KeyringStore._access_token_key)
             _keyring.delete_password(for_endpoint, KeyringStore._refresh_token_key)
+            try:
+                _keyring.delete_password(for_endpoint, KeyringStore._id_token_key)
+            except PasswordDeleteError as e:
+                logging.debug(f"Id token not found in key store, not deleting. Error: {e}")
         except NoKeyringError as e:
             logging.debug(f"KeyRing not available, tokens will not be cached. Error: {e}")