Factor metadata interface into a separate module

oauth2client/contrib/gce.py

@@ -21,25 +21,20 @@
 import logging
 import warnings
 
-import httplib2
-from six.moves import http_client
-from six.moves import urllib
 
 from oauth2client._helpers import _from_bytes
 from oauth2client import util
-from oauth2client.client import HttpAccessTokenRefreshError
 from oauth2client.client import AssertionCredentials
+from oauth2client.contrib.metadata import MetadataServer
 
 
 __author__ = '[email protected] (Joe Gregorio)'
 
 logger = logging.getLogger(__name__)
 
-# URI Template for the endpoint that returns access_tokens.
-_METADATA_ROOT = ('http://metadata.google.internal/computeMetadata/v1/'
-                  'instance/service-accounts/default/')
-META = _METADATA_ROOT + 'token'
-_DEFAULT_EMAIL_METADATA = _METADATA_ROOT + 'email'
+# Backwards Compat
+META = ('http://metadata.google.internal/computeMetadata/v1/'
+        'instance/service-accounts/default/token')
 _SCOPES_WARNING = """\
 You have requested explicit scopes to be used with a GCE service account.
 Using this argument will have no effect on the actual scopes for tokens
@@ -48,30 +43,6 @@
 """
 
 
-def _get_service_account_email(http_request=None):
-    """Get the GCE service account email from the current environment.
-
-    Args:
-        http_request: callable, (Optional) a callable that matches the method
-                      signature of httplib2.Http.request, used to make
-                      the request to the metadata service.
-
-    Returns:
-        tuple, A pair where the first entry is an optional response (from a
-        failed request) and the second is service account email found (as
-        a string).
-    """
-    if http_request is None:
-        http_request = httplib2.Http().request
-    response, content = http_request(
-        _DEFAULT_EMAIL_METADATA, headers={'Metadata-Flavor': 'Google'})
-    if response.status == http_client.OK:
-        content = _from_bytes(content)
-        return None, content
-    else:
-        return response, content
-
-
 class AppAssertionCredentials(AssertionCredentials):
     """Credentials object for Compute Engine Assertion Grants
 
@@ -85,8 +56,11 @@ class AppAssertionCredentials(AssertionCredentials):
     information to generate and refresh its own access tokens.
     """
 
+    NON_SERIALIZED_MEMBERS = AssertionCredentials.NON_SERIALIZED_MEMBERS.extend(
+        ['_metadata', 'kwargs'])
+
     @util.positional(2)
-    def __init__(self, scope='', **kwargs):
+    def __init__(self, scope='', metadata_server=None, **kwargs):
         """Constructor for AppAssertionCredentials
 
         Args:
@@ -102,10 +76,11 @@ def __init__(self, scope='', **kwargs):
         self.scope = util.scopes_to_string(scope)
         self.kwargs = kwargs
 
+        self._metadata = metadata_server or MetadataServer()
+
         # Assertion type is no longer used, but still in the
         # parent class signature.
         super(AppAssertionCredentials, self).__init__(None)
-        self._service_account_email = None
 
     @classmethod
     def from_json(cls, json_data):
@@ -125,21 +100,8 @@ def _refresh(self, http_request):
         Raises:
             HttpAccessTokenRefreshError: When the refresh fails.
         """
-        response, content = http_request(
-            META, headers={'Metadata-Flavor': 'Google'})
-        content = _from_bytes(content)
-        if response.status == http_client.OK:
-            try:
-                token_content = json.loads(content)
-            except Exception as e:
-                raise HttpAccessTokenRefreshError(str(e),
-                                                  status=response.status)
-            self.access_token = token_content['access_token']
-        else:
-            if response.status == http_client.NOT_FOUND:
-                content += (' This can occur if a VM was created'
-                            ' with no service account or scopes.')
-            raise HttpAccessTokenRefreshError(content, status=response.status)
+        self.access_token, self.token_expiry = self._metadata.get_token(
+            http_request=http_request)
 
     @property
     def serialization_data(self):
@@ -183,12 +145,4 @@ def service_account_email(self):
             AttributeError, if the email can not be retrieved from the Google
             Compute Engine metadata service.
         """
-        if self._service_account_email is None:
-            failure, email = _get_service_account_email()
-            if failure is None:
-                self._service_account_email = email
-            else:
-                raise AttributeError('Failed to retrieve the email from the '
-                                     'Google Compute Engine metadata service',
-                                     failure, email)
-        return self._service_account_email
+        return self._metadata.get_service_account_info()['email']

oauth2client/contrib/metadata.py

@@ -0,0 +1,144 @@
+# Copyright 2014 Google Inc. All rights reserved.
+#
+# 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.
+
+"""Thin wrapper class for talking to the GCE Metadata Server."""
+import datetime
+import httplib2
+import json
+
+from six.moves import http_client
+
+from oauth2client._helpers import _from_bytes
+from oauth2client.client import _UTCNOW
+from oauth2client.client import HttpAccessTokenRefreshError
+
+class NestedDict(dict):
+    """Stores a dict and allows setting and retrieving
+     values by path (list of keys)."""
+
+    def get_path(self, path):
+        leaf = self
+        for key in path:
+            leaf = leaf.get(key)
+            if leaf is None:
+                return None
+        return leaf
+
+    def set_path(self, path, value):
+        leaf = self
+        for key in path[:-1]:
+            leaf = leaf.setdefault(key, {})
+        leaf[path[-1]] = value
+
+
+class MetadataServerHttpError(Exception):
+    """Error for Http failures originating from the Metadata Server"""
+
+
+class MetadataServer:
+    """handles requests to and from the metadata server,
+     and caches requests by default"""
+
+    def __init__(self,
+                 client=None,
+                 cache=None,
+                 root='http://metadata.google.internal/computeMetadata/v1/'):
+        self._client = client or httplib2.Http()
+        self._root = root
+        self.cache = cache or NestedDict()
+
+    def _make_request(self, path, recursive=True, http_request=None):
+        if path is None:
+            path = []
+
+        if not http_request:
+            http_request = self._client.request
+
+        r_string = '/?recursive=true' if recursive else ''
+        full_path = self._root + '/'.join(path) + r_string
+        response, content = http_request(
+            full_path,
+            headers={'Metadata-Flavor': 'Google'}
+        )
+        if response.status == http_client.OK:
+            decoded = _from_bytes(content)
+            if response['content-type'] == 'application/json':
+                return json.loads(decoded)
+            else:
+                return decoded
+        else:
+            msg = (
+                'Failed to retrieve {path} from the Google Compute Engine'
+                'metadata service. Response:\n{error}'
+            ).format(path=full_path, error=response)
+            raise MetadataServerHttpError(msg)
+
+    def get(self, path, use_cache=True, recursive=True, http_request=None):
+        """ Retrieve a value from the metadata server.
+        :param path: Path on the metadata server to fetch from
+        :param use_cache: Use a cached value (if available) and update the cache (if not)
+        :param recursive: True if this is not a leaf
+        :param http_request: callable, a callable that matches the method
+            signature of httplib2.Http.request, used to make
+            the refresh request.
+        :return: The value from the metadata server (String if recursive=False, dict otherwise)
+        """
+
+        if use_cache:
+            cached_value = self.cache.get_path(path)
+            if cached_value is not None:
+                return cached_value
+        value = self._make_request(path, recursive=recursive, http_request=http_request)
+        if use_cache:
+            self.cache.set_path(path, value)
+        return value
+
+    def get_service_account_info(self, service_account='default', http_request=None):
+        """ Get information about a service account from the metadata server.
+        :param service_account: a service account email. Left blank information for
+            the default service account of current compute engine instance will be looked up.
+        :param http_request: callable, a callable that matches the method
+            signature of httplib2.Http.request, used to make
+            the refresh request.
+        :return: A dictionary with information about the specified service account.
+        """
+        return self.get(
+            ['instance', 'service-accounts', service_account],
+            use_cache=True,
+            recursive=True,
+            http_request=http_request
+        )
+
+    def get_token(self, service_account='default', http_request=None):
+        """Fetch an OAuth access token from the metadata server
+        :param service_account: a service account email. Left blank information for
+            the default service account of current compute engine instance will be looked up.
+        :param http_request: callable, a callable that matches the method
+            signature of httplib2.Http.request, used to make
+            the refresh request.
+        :return:
+        """
+        try:
+            token_json = self.get(
+                ['instance', 'service-accounts', service_account, 'token'],
+                use_cache=False,
+                recursive=False,
+                http_request=http_request
+            )
+        except MetadataServerHttpError as failed_fetch:
+            raise HttpAccessTokenRefreshError(str(failed_fetch))
+
+        token_expiry = _UTCNOW() + datetime.timedelta(
+            seconds=token_json['expires_in'])
+        return token_json['access_token'], token_expiry