diff --git a/sdk/python/build.sh b/sdk/python/build.sh index 50609cceed7..e901992e05e 100755 --- a/sdk/python/build.sh +++ b/sdk/python/build.sh @@ -23,7 +23,7 @@ # Setup: # apt-get update -y # apt-get install --no-install-recommends -y -q default-jdk -# wget http://central.maven.org/maven2/io/swagger/swagger-codegen-cli/2.3.1/swagger-codegen-cli-2.3.1.jar -O /tmp/swagger-codegen-cli.jar +# wget http://central.maven.org/maven2/io/swagger/swagger-codegen-cli/2.4.1/swagger-codegen-cli-2.4.1.jar -O /tmp/swagger-codegen-cli.jar get_abs_filename() { # $1 : relative filename diff --git a/sdk/python/kfp/_auth.py b/sdk/python/kfp/_auth.py new file mode 100644 index 00000000000..6ecfffccf64 --- /dev/null +++ b/sdk/python/kfp/_auth.py @@ -0,0 +1,110 @@ +# Copyright 2018 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import google.auth +import google.auth.app_engine +import google.auth.compute_engine.credentials +import google.auth.iam +from google.auth.transport.requests import Request +import google.oauth2.credentials +import google.oauth2.service_account +import requests_toolbelt.adapters.appengine + +IAM_SCOPE = 'https://www.googleapis.com/auth/iam' +OAUTH_TOKEN_URI = 'https://www.googleapis.com/oauth2/v4/token' + +def get_auth_token(client_id): + """Gets auth token from default service account. + + If no service account credential is found, returns None. + """ + service_account_credentials = get_service_account_credentials(client_id) + if service_account_credentials: + return get_google_open_id_connect_token(service_account_credentials) + return None + +def get_service_account_credentials(client_id): + # Figure out what environment we're running in and get some preliminary + # information about the service account. + bootstrap_credentials, _ = google.auth.default( + scopes=[IAM_SCOPE]) + if isinstance(bootstrap_credentials, + google.oauth2.credentials.Credentials): + logging.info('Found OAuth2 credentials and skip SA auth.') + return None + elif isinstance(bootstrap_credentials, + google.auth.app_engine.Credentials): + requests_toolbelt.adapters.appengine.monkeypatch() + + # For service account's using the Compute Engine metadata service, + # service_account_email isn't available until refresh is called. + bootstrap_credentials.refresh(Request()) + signer_email = bootstrap_credentials.service_account_email + if isinstance(bootstrap_credentials, + google.auth.compute_engine.credentials.Credentials): + # Since the Compute Engine metadata service doesn't expose the service + # account key, we use the IAM signBlob API to sign instead. + # In order for this to work: + # + # 1. Your VM needs the https://www.googleapis.com/auth/iam scope. + # You can specify this specific scope when creating a VM + # through the API or gcloud. When using Cloud Console, + # you'll need to specify the "full access to all Cloud APIs" + # scope. A VM's scopes can only be specified at creation time. + # + # 2. The VM's default service account needs the "Service Account Actor" + # role. This can be found under the "Project" category in Cloud + # Console, or roles/iam.serviceAccountActor in gcloud. + signer = google.auth.iam.Signer( + Request(), bootstrap_credentials, signer_email) + else: + # A Signer object can sign a JWT using the service account's key. + signer = bootstrap_credentials.signer + + # Construct OAuth 2.0 service account credentials using the signer + # and email acquired from the bootstrap credentials. + return google.oauth2.service_account.Credentials( + signer, signer_email, token_uri=OAUTH_TOKEN_URI, additional_claims={ + 'target_audience': client_id + }) + +def get_google_open_id_connect_token(service_account_credentials): + """Get an OpenID Connect token issued by Google for the service account. + This function: + 1. Generates a JWT signed with the service account's private key + containing a special "target_audience" claim. + 2. Sends it to the OAUTH_TOKEN_URI endpoint. Because the JWT in #1 + has a target_audience claim, that endpoint will respond with + an OpenID Connect token for the service account -- in other words, + a JWT signed by *Google*. The aud claim in this JWT will be + set to the value from the target_audience claim in #1. + For more information, see + https://developers.google.com/identity/protocols/OAuth2ServiceAccount . + The HTTP/REST example on that page describes the JWT structure and + demonstrates how to call the token endpoint. (The example on that page + shows how to get an OAuth2 access token; this code is using a + modified version of it to get an OpenID Connect token.) + """ + + service_account_jwt = ( + service_account_credentials._make_authorization_grant_assertion()) + request = google.auth.transport.requests.Request() + body = { + 'assertion': service_account_jwt, + 'grant_type': google.oauth2._client._JWT_GRANT_TYPE, + } + token_response = google.oauth2._client._token_endpoint_request( + request, OAUTH_TOKEN_URI, body) + return token_response['id_token'] \ No newline at end of file diff --git a/sdk/python/kfp/_client.py b/sdk/python/kfp/_client.py index a1111a29d3f..ccd91b8427a 100644 --- a/sdk/python/kfp/_client.py +++ b/sdk/python/kfp/_client.py @@ -24,6 +24,7 @@ from .compiler import compiler from .compiler import _k8s_helper +from ._auth import get_auth_token class Client(object): """ API Client for KubeFlow Pipeline. @@ -32,7 +33,7 @@ class Client(object): # in-cluster DNS name of the pipeline service IN_CLUSTER_DNS_NAME = 'ml-pipeline.kubeflow.svc.cluster.local:8888' - def __init__(self, host=None): + def __init__(self, host=None, client_id=None): """Create a new instance of kfp client. Args: @@ -41,6 +42,7 @@ def __init__(self, host=None): in the same cluster (such as a Jupyter instance spawned by Kubeflow's JupyterHub). If you have a different connection to cluster, such as a kubectl proxy connection, then set it to something like "127.0.0.1:8080/pipeline". + client_id: The client ID used by Identity-Aware Proxy. """ try: @@ -55,17 +57,28 @@ def __init__(self, host=None): self._host = host + token = None + if host and client_id: + token = get_auth_token(client_id) + config = kfp_run.configuration.Configuration() config.host = host if host else Client.IN_CLUSTER_DNS_NAME + self._configure_auth(config, token) api_client = kfp_run.api_client.ApiClient(config) self._run_api = kfp_run.api.run_service_api.RunServiceApi(api_client) config = kfp_experiment.configuration.Configuration() config.host = host if host else Client.IN_CLUSTER_DNS_NAME + self._configure_auth(config, token) api_client = kfp_experiment.api_client.ApiClient(config) self._experiment_api = \ kfp_experiment.api.experiment_service_api.ExperimentServiceApi(api_client) + def _configure_auth(self, config, token): + if token: + config.api_key['authorization'] = token + config.api_key_prefix['authorization'] = 'Bearer' + def _is_ipython(self): """Returns whether we are running in notebook.""" try: diff --git a/sdk/python/setup.py b/sdk/python/setup.py index bbca8a7b94b..c8753d45c46 100644 --- a/sdk/python/setup.py +++ b/sdk/python/setup.py @@ -20,7 +20,8 @@ VERSION = '0.1' REQUIRES = ['urllib3 >= 1.15', 'six >= 1.10', 'certifi', 'python-dateutil', 'PyYAML', - 'google-cloud-storage == 1.13.0', 'kubernetes == 8.0.0'] + 'google-cloud-storage == 1.13.0', 'kubernetes == 8.0.0', 'PyJWT==1.6.4', + 'cryptography==2.4.2', 'google-auth==1.6.1', 'requests_toolbelt==0.8.0'] setup( name=NAME,