Add proper GCP config loader and refresher

Author: mbohlool

config/kube_config.py

@@ -14,17 +14,22 @@
 
 import atexit
 import base64
+import datetime
 import os
 import tempfile
+import time
 
+import google.auth
+import google.auth.transport.requests
 import urllib3
 import yaml
-from google.oauth2.credentials import Credentials
 
 from kubernetes.client import ApiClient, ConfigurationObject, configuration
 
 from .config_exception import ConfigException
+from .rfc3339 import tf_from_timestamp, timestamp_from_tf
 
+EXPIRY_SKEW_PREVENTION_DELAY_S = 600
 KUBE_CONFIG_DEFAULT_LOCATION = os.environ.get('KUBECONFIG', '~/.kube/config')
 _temp_files = {}
 
@@ -54,6 +59,17 @@ def _create_temp_file_with_content(content):
     return name
 
 
+def _is_expired(expiry):
+    tf = tf_from_timestamp(expiry)
+    n = time.time()
+    return tf + EXPIRY_SKEW_PREVENTION_DELAY_S <= n
+
+
+def _datetime_to_rfc3339(dt):
+    tf = (dt - datetime.datetime.utcfromtimestamp(0)).total_seconds()
+    return timestamp_from_tf(tf, time_offset="Z")
+
+
 class FileOrData(object):
     """Utility class to read content of obj[%data_key_name] or file's
      content of obj[%file_key_name] and represent it as file or data.
@@ -110,19 +126,26 @@ class KubeConfigLoader(object):
     def __init__(self, config_dict, active_context=None,
                  get_google_credentials=None,
                  client_configuration=configuration,
-                 config_base_path=""):
+                 config_base_path="",
+                 config_persister=None):
         self._config = ConfigNode('kube-config', config_dict)
         self._current_context = None
         self._user = None
         self._cluster = None
         self.set_active_context(active_context)
         self._config_base_path = config_base_path
+        self._config_persister = config_persister
+
+        def _refresh_credentials():
+            credentials, project_id = google.auth.default()
+            request = google.auth.transport.requests.Request()
+            credentials.refresh(request)
+            return credentials
+
         if get_google_credentials:
             self._get_google_credentials = get_google_credentials
         else:
-            self._get_google_credentials = lambda: (
-                GoogleCredentials.get_application_default()
-                .get_access_token().access_token)
+            self._get_google_credentials = _refresh_credentials
         self._client_configuration = client_configuration
 
     def set_active_context(self, context_name=None):
@@ -166,16 +189,32 @@ def _load_authentication(self):
     def _load_gcp_token(self):
         if 'auth-provider' not in self._user:
             return
-        if 'name' not in self._user['auth-provider']:
+        provider = self._user['auth-provider']
+        if 'name' not in provider:
             return
-        if self._user['auth-provider']['name'] != 'gcp':
+        if provider['name'] != 'gcp':
             return
-        # Ignore configs in auth-provider and rely on GoogleCredentials
-        # caching and refresh mechanism.
-        # TODO: support gcp command based token ("cmd-path" config).
-        self.token = "Bearer %s" % self._get_google_credentials()
+
+        if (('config' not in provider) or
+                ('access-token' not in provider['config']) or
+                ('expiry' in provider['config'] and
+                 _is_expired(provider['config']['expiry']))):
+            # token is not available or expired, refresh it
+            self._refresh_gcp_token()
+
+        self.token = "Bearer %s" % provider['config']['access-token']
         return self.token
 
+    def _refresh_gcp_token(self):
+        if 'config' not in self._user['auth-provider']:
+            self._user['auth-provider'].value['config'] = {}
+        provider = self._user['auth-provider']['config']
+        credentials = self._get_google_credentials()
+        provider.value['access-token'] = credentials.token
+        provider.value['expiry'] = _datetime_to_rfc3339(credentials.expiry)
+        if self._config_persister:
+            self._config_persister(self._config.value)
+
     def _load_user_token(self):
         token = FileOrData(
             self._user, 'tokenFile', 'token',
@@ -289,6 +328,11 @@ def _get_kube_config_loader_for_yaml_file(filename, **kwargs):
             **kwargs)
 
 
+def _save_kube_config(filename, config_map):
+    with open(filename, 'w') as f:
+        yaml.safe_dump(config_map, f, default_flow_style=False)
+
+
 def list_kube_config_contexts(config_file=None):
 
     if config_file is None:
@@ -299,7 +343,8 @@ def list_kube_config_contexts(config_file=None):
 
 
 def load_kube_config(config_file=None, context=None,
-                     client_configuration=configuration):
+                     client_configuration=configuration,
+                     persist_config=True):
     """Loads authentication and cluster information from kube-config file
     and stores them in kubernetes.client.configuration.
 
@@ -308,21 +353,32 @@ def load_kube_config(config_file=None, context=None,
         from config file will be used.
     :param client_configuration: The kubernetes.client.ConfigurationObject to
         set configs to.
+    :param persist_config: If True and config changed (e.g. GCP token refresh)
+        the provided config file will be updated.
     """
 
     if config_file is None:
         config_file = os.path.expanduser(KUBE_CONFIG_DEFAULT_LOCATION)
 
+    config_persister = None
+    if persist_config:
+        config_persister = lambda config_map, config_file=config_file: (
+            _save_kube_config(config_file, config_map))
     _get_kube_config_loader_for_yaml_file(
         config_file, active_context=context,
-        client_configuration=client_configuration).load_and_set()
+        client_configuration=client_configuration,
+        config_persister=config_persister).load_and_set()
 
 
-def new_client_from_config(config_file=None, context=None):
+def new_client_from_config(
+        config_file=None,
+        context=None,
+        persist_config=True):
     """Loads configuration the same as load_kube_config but returns an ApiClient
     to be used with any API object. This will allow the caller to concurrently
     talk with multiple clusters."""
     client_config = ConfigurationObject()
     load_kube_config(config_file=config_file, context=context,
-                     client_configuration=client_config)
+                     client_configuration=client_config,
+                     persist_config=persist_config)
     return ApiClient(config=client_config)

config/kube_config_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import base64
+import datetime
 import os
 import shutil
 import tempfile
@@ -26,6 +27,7 @@
                           _cleanup_temp_files, _create_temp_file_with_content,
                           list_kube_config_contexts, load_kube_config,
                           new_client_from_config)
+from .rfc3339 import timestamp_from_tf
 
 BEARER_TOKEN_FORMAT = "Bearer %s"
 
@@ -304,6 +306,13 @@ class TestKubeConfigLoader(BaseTestCase):
                     "user": "gcp"
                 }
             },
+            {
+                "name": "expired_gcp",
+                "context": {
+                    "cluster": "default",
+                    "user": "expired_gcp"
+                }
+            },
             {
                 "name": "user_pass",
                 "context": {
@@ -397,7 +406,24 @@ class TestKubeConfigLoader(BaseTestCase):
                 "user": {
                     "auth-provider": {
                         "name": "gcp",
-                        "access_token": "not_used",
+                        "config": {
+                            "access-token": TEST_DATA_BASE64,
+                        }
+                    },
+                    "token": TEST_DATA_BASE64,  # should be ignored
+                    "username": TEST_USERNAME,  # should be ignored
+                    "password": TEST_PASSWORD,  # should be ignored
+                }
+            },
+            {
+                "name": "expired_gcp",
+                "user": {
+                    "auth-provider": {
+                        "name": "gcp",
+                        "config": {
+                            "access-token": TEST_DATA_BASE64,
+                            "expiry": timestamp_from_tf(0),
+                        }
                     },
                     "token": TEST_DATA_BASE64,  # should be ignored
                     "username": TEST_USERNAME,  # should be ignored
@@ -464,24 +490,38 @@ def test_load_user_token(self):
         self.assertTrue(loader._load_user_token())
         self.assertEqual(BEARER_TOKEN_FORMAT % TEST_DATA_BASE64, loader.token)
 
-    def test_gcp(self):
+    def test_gcp_no_refresh(self):
         expected = FakeConfig(
             host=TEST_HOST,
-            token=BEARER_TOKEN_FORMAT % TEST_ANOTHER_DATA_BASE64)
+            token=BEARER_TOKEN_FORMAT % TEST_DATA_BASE64)
         actual = FakeConfig()
         KubeConfigLoader(
             config_dict=self.TEST_KUBE_CONFIG,
             active_context="gcp",
             client_configuration=actual,
-            get_google_credentials=lambda: TEST_ANOTHER_DATA_BASE64) \
+            get_google_credentials=lambda: "SHOULD NOT BE CALLED") \
             .load_and_set()
         self.assertEqual(expected, actual)
 
-    def test_load_gcp_token(self):
+    def test_load_gcp_token_no_refresh(self):
         loader = KubeConfigLoader(
             config_dict=self.TEST_KUBE_CONFIG,
             active_context="gcp",
-            get_google_credentials=lambda: TEST_ANOTHER_DATA_BASE64)
+            get_google_credentials=lambda: "SHOULD NOT BE CALLED")
+        self.assertTrue(loader._load_gcp_token())
+        self.assertEqual(BEARER_TOKEN_FORMAT % TEST_DATA_BASE64,
+                         loader.token)
+
+    def test_load_gcp_token_with_refresh(self):
+
+        def cred(): return None
+        cred.token = TEST_ANOTHER_DATA_BASE64
+        cred.expiry = datetime.datetime.now()
+
+        loader = KubeConfigLoader(
+            config_dict=self.TEST_KUBE_CONFIG,
+            active_context="expired_gcp",
+            get_google_credentials=lambda: cred)
         self.assertTrue(loader._load_gcp_token())
         self.assertEqual(BEARER_TOKEN_FORMAT % TEST_ANOTHER_DATA_BASE64,
                          loader.token)

config/rfc3339.MD

@@ -0,0 +1 @@
+The (rfc3339.py)[rfc3339.py] file is copied from [this site](http://home.blarg.net/~steveha/pyfeed.html) because PyFeed is not available in PyPi.