Move compute engine metadata interface into a separate module (#520)

oauth2client/contrib/_metadata.py

@@ -0,0 +1,126 @@
+# Copyright 2016 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.
+
+"""Provides helper methods for talking to the Compute Engine metadata server.
+
+See https://cloud.google.com/compute/docs/metadata
+"""
+
+import datetime
+import httplib2
+import json
+
+from six.moves import http_client
+from six.moves.urllib import parse as urlparse
+
+from oauth2client._helpers import _from_bytes
+from oauth2client.client import _UTCNOW
+from oauth2client import util
+
+
+METADATA_ROOT = 'http://metadata.google.internal/computeMetadata/v1/'
+METADATA_HEADERS = {'Metadata-Flavor': 'Google'}
+
+
+def get(path, http_request=None, root=METADATA_ROOT, recursive=None):
+    """Fetch a resource from the metadata server.
+
+    Args:
+        path: A string indicating the resource to retrieve. For example,
+            'instance/service-accounts/defualt'
+        http_request: A callable that matches the method
+            signature of httplib2.Http.request. Used to make the request to the
+            metadataserver.
+        root: A string indicating the full path to the metadata server root.
+        recursive: A boolean indicating whether to do a recursive query of
+            metadata. See
+            https://cloud.google.com/compute/docs/metadata#aggcontents
+
+    Returns:
+        A dictionary if the metadata server returns JSON, otherwise a string.
+
+    Raises:
+        httplib2.Httplib2Error if an error corrured while retrieving metadata.
+    """
+    if not http_request:
+        http_request = httplib2.Http().request
+
+    url = urlparse.urljoin(root, path)
+    url = util._add_query_parameter(url, 'recursive', recursive)
+
+    response, content = http_request(
+        url,
+        headers=METADATA_HEADERS
+    )
+
+    if response.status == http_client.OK:
+        decoded = _from_bytes(content)
+        if response['content-type'] == 'application/json':
+            return json.loads(decoded)
+        else:
+            return decoded
+    else:
+        raise httplib2.HttpLib2Error(
+            'Failed to retrieve {0} from the Google Compute Engine'
+            'metadata service. Response:\n{1}'.format(url, response))
+
+
+def get_service_account_info(service_account='default', http_request=None):
+    """Get information about a service account from the metadata server.
+
+    Args:
+        service_account: An email specifying the service account for which to
+            look up information. Default will be information for the "default"
+            service account of the current compute engine instance.
+        http_request: A callable that matches the method
+            signature of httplib2.Http.request. Used to make the request to the
+            metadata server.
+    Returns:
+         A dictionary with information about the specified service account,
+         for example:
+
+            {
+                'email': '...',
+                'scopes': ['scope', ...],
+                'aliases': ['default', '...']
+            }
+    """
+    return get(
+        'instance/service-accounts/{0}'.format(service_account),
+        recursive=True,
+        http_request=http_request)
+
+
+def get_token(service_account='default', http_request=None):
+    """Fetch an oauth token for the
+
+    Args:
+        service_account: An email specifying the service account this token
+            should represent. Default will be a token for the "default" service
+            account of the current compute engine instance.
+        http_request: A callable that matches the method
+            signature of httplib2.Http.request. Used to make the request to the
+            metadataserver.
+
+    Returns:
+         A tuple of (access token, token expiration), where access token is the
+         access token as a string and token expiration is a datetime object
+         that indicates when the access token will expire.
+    """
+    token_json = get(
+        'instance/service-accounts/{0}/token'.format(service_account),
+        http_request=http_request)
+    token_expiry = _UTCNOW() + datetime.timedelta(
+        seconds=token_json['expires_in'])
+    return token_json['access_token'], token_expiry

oauth2client/contrib/gce.py

@@ -17,30 +17,23 @@
 Utilities for making it easier to use OAuth 2.0 on Google Compute Engine.
 """
 
-import datetime
 import json
 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.client import HttpAccessTokenRefreshError
+from oauth2client.contrib import _metadata
 
 
 __author__ = 'jcgregorio@google.com (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'
 _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
@@ -49,30 +42,6 @@ can't be overridden in the request.
 """
 
 
-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
 
@@ -106,6 +75,8 @@ class AppAssertionCredentials(AssertionCredentials):
         # Assertion type is no longer used, but still in the
         # parent class signature.
         super(AppAssertionCredentials, self).__init__(None)
+
+        # Cache until Metadata Server supports Cache-Control Header
         self._service_account_email = None
 
     @classmethod
@@ -126,23 +97,11 @@ class AppAssertionCredentials(AssertionCredentials):
         Raises:
             HttpAccessTokenRefreshError: When the refresh fails.
         """
-        response, content = http_request(
+        try:
-            META, headers={'Metadata-Flavor': 'Google'})
+            self.access_token, self.token_expiry = _metadata.get_token(
-        content = _from_bytes(content)
+                http_request=http_request)
-        if response.status == http_client.OK:
+        except httplib2.HttpLib2Error as e:
-            try:
+            raise HttpAccessTokenRefreshError(str(e))
-                token_content = json.loads(content)
-            except Exception as e:
-                raise HttpAccessTokenRefreshError(str(e),
-                                                  status=response.status)
-            self.access_token = token_content['access_token']
-            delta = datetime.timedelta(seconds=int(token_content['expires_in']))
-            self.token_expiry = delta + datetime.datetime.utcnow()
-        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)
 
     @property
     def serialization_data(self):
@@ -187,11 +146,6 @@ class AppAssertionCredentials(AssertionCredentials):
             Compute Engine metadata service.
         """
         if self._service_account_email is None:
-            failure, email = _get_service_account_email()
+            self._service_account_email = (
-            if failure is None:
+                _metadata.get_service_account_info()['email'])
-                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

tests/contrib/test_gce.py

@@ -14,24 +14,20 @@
 
 """Unit tests for oauth2client.contrib.gce."""
 
+import datetime
 import json
-from datetime import datetime
+
+import mock
 from six.moves import http_client
 from six.moves import urllib
 import unittest2
 
-import mock
-
-import httplib2
-from oauth2client._helpers import _to_bytes
-from oauth2client.client import AccessTokenRefreshError
 from oauth2client.client import Credentials
 from oauth2client.client import save_to_well_known_file
-from oauth2client.contrib.gce import _DEFAULT_EMAIL_METADATA
+from oauth2client.client import HttpAccessTokenRefreshError
-from oauth2client.contrib.gce import _get_service_account_email
 from oauth2client.contrib.gce import _SCOPES_WARNING
 from oauth2client.contrib.gce import AppAssertionCredentials
-
+from tests.contrib.test_metadata import request_mock
 
 __author__ = 'jcgregorio@google.com (Joe Gregorio)'
 
@@ -61,80 +57,29 @@ class AppAssertionCredentialsTests(unittest2.TestCase):
         self.assertEqual(credentials.access_token,
                          credentials_from_json.access_token)
 
-    def _refresh_success_helper(self, bytes_response=False):
+    @mock.patch('oauth2client.contrib._metadata.get_token',
-        access_token = u'this-is-a-token'
+                side_effect=[('A', datetime.datetime.min),
-        expires_in = 600
+                             ('B', datetime.datetime.max)])
-        return_val = json.dumps({
+    def test_refresh_token(self, metadata):
-            u'access_token': access_token,
-            u'expires_in': expires_in
-        })
-        if bytes_response:
-            return_val = _to_bytes(return_val)
-        http = mock.MagicMock()
-        http.request = mock.MagicMock(
-            return_value=(mock.Mock(status=http_client.OK), return_val))
-
         credentials = AppAssertionCredentials()
-        self.assertEquals(None, credentials.access_token)
+        self.assertIsNone(credentials.access_token)
-        credentials.refresh(http)
+        credentials.get_access_token()
-        self.assertEquals(access_token, credentials.access_token)
+        self.assertEqual(credentials.access_token, 'A')
+        self.assertTrue(credentials.access_token_expired)
+        credentials.get_access_token()
+        self.assertEqual(credentials.access_token, 'B')
         self.assertFalse(credentials.access_token_expired)
-        self.assertTrue(credentials.token_expiry > datetime.utcnow())
-
-        base_metadata_uri = (
-            'http://metadata.google.internal/computeMetadata/v1/instance/'
-            'service-accounts/default/token')
-        http.request.assert_called_once_with(
-            base_metadata_uri, headers={'Metadata-Flavor': 'Google'})
-
-    def test_refresh_success(self):
-        self._refresh_success_helper(bytes_response=False)
-
-    def test_refresh_success_bytes(self):
-        self._refresh_success_helper(bytes_response=True)
-
-    def test_refresh_failure_bad_json(self):
-        http = mock.MagicMock()
-        content = '{BADJSON'
-        http.request = mock.MagicMock(
-            return_value=(mock.Mock(status=http_client.OK), content))
 
+    def test_refresh_token_failed_fetch(self):
+        http_request = request_mock(
+            http_client.NOT_FOUND,
+            'application/json',
+            json.dumps({'access_token': 'a', 'expires_in': 100})
+        )
         credentials = AppAssertionCredentials()
-        self.assertRaises(AccessTokenRefreshError, credentials.refresh, http)
 
-    def test_refresh_failure_400(self):
+        with self.assertRaises(HttpAccessTokenRefreshError):
-        http = mock.MagicMock()
+            credentials._refresh(http_request=http_request)
-        content = '{}'
-        http.request = mock.MagicMock(
-            return_value=(mock.Mock(status=http_client.BAD_REQUEST), content))
-
-        credentials = AppAssertionCredentials()
-        exception_caught = None
-        try:
-            credentials.refresh(http)
-        except AccessTokenRefreshError as exc:
-            exception_caught = exc
-
-        self.assertNotEqual(exception_caught, None)
-        self.assertEqual(str(exception_caught), content)
-
-    def test_refresh_failure_404(self):
-        http = mock.MagicMock()
-        content = '{}'
-        http.request = mock.MagicMock(
-            return_value=(mock.Mock(status=http_client.NOT_FOUND), content))
-
-        credentials = AppAssertionCredentials()
-        exception_caught = None
-        try:
-            credentials.refresh(http)
-        except AccessTokenRefreshError as exc:
-            exception_caught = exc
-
-        self.assertNotEqual(exception_caught, None)
-        expanded_content = content + (' This can occur if a VM was created'
-                                      ' with no service account or scopes.')
-        self.assertEqual(str(exception_caught), expanded_content)
 
     def test_serialization_data(self):
         credentials = AppAssertionCredentials()
@@ -165,60 +110,13 @@ class AppAssertionCredentialsTests(unittest2.TestCase):
         with self.assertRaises(NotImplementedError):
             credentials.sign_blob(b'blob')
 
-    @mock.patch('oauth2client.contrib.gce._get_service_account_email',
+    @mock.patch('oauth2client.contrib._metadata.get_service_account_info',
-                return_value=(None, 'retrieved@email.com'))
+                return_value={'email': 'a@example.com'})
-    def test_service_account_email(self, get_email):
+    def test_service_account_email(self, metadata):
-        credentials = AppAssertionCredentials([])
-        self.assertIsNone(credentials._service_account_email)
-        self.assertEqual(credentials.service_account_email,
-                         get_email.return_value[1])
-        self.assertIsNotNone(credentials._service_account_email)
-        get_email.assert_called_once_with()
-
-    @mock.patch('oauth2client.contrib.gce._get_service_account_email')
-    def test_service_account_email_already_set(self, get_email):
-        credentials = AppAssertionCredentials([])
-        acct_name = 'existing@email.com'
-        credentials._service_account_email = acct_name
-        self.assertEqual(credentials.service_account_email, acct_name)
-        get_email.assert_not_called()
-
-    @mock.patch('oauth2client.contrib.gce._get_service_account_email')
-    def test_service_account_email_failure(self, get_email):
-        # Set-up the mock.
-        bad_response = httplib2.Response({'status': http_client.NOT_FOUND})
-        content = b'bad-bytes-nothing-here'
-        get_email.return_value = (bad_response, content)
-        # Test the failure.
-        credentials = AppAssertionCredentials([])
-        self.assertIsNone(credentials._service_account_email)
-        with self.assertRaises(AttributeError) as exc_manager:
-            getattr(credentials, 'service_account_email')
-
-        error_msg = ('Failed to retrieve the email from the '
-                     'Google Compute Engine metadata service')
-        self.assertEqual(
-            exc_manager.exception.args,
-            (error_msg, bad_response, content))
-        self.assertIsNone(credentials._service_account_email)
-        get_email.assert_called_once_with()
-
-    def test_get_access_token(self):
-        http = mock.MagicMock()
-        http.request = mock.MagicMock(
-            return_value=(mock.Mock(status=http_client.OK),
-                          '{"access_token": "this-is-a-token", '
-                          '"expires_in": 600}'))
-
         credentials = AppAssertionCredentials()
-        token = credentials.get_access_token(http=http)
+        # Assert that service account isn't pre-fetched
-        self.assertEqual('this-is-a-token', token.access_token)
+        metadata.assert_not_called()
-        self.assertGreaterEqual(600, token.expires_in)
+        self.assertEqual(credentials.service_account_email, 'a@example.com')
-
-        http.request.assert_called_once_with(
-            'http://metadata.google.internal/computeMetadata/v1/instance/'
-            'service-accounts/default/token',
-            headers={'Metadata-Flavor': 'Google'})
 
     def test_save_to_well_known_file(self):
         import os
@@ -232,43 +130,5 @@ class AppAssertionCredentialsTests(unittest2.TestCase):
             os.path.isdir = ORIGINAL_ISDIR
 
 
-class Test__get_service_account_email(unittest2.TestCase):
-
-    def test_success(self):
-        http_request = mock.MagicMock()
-        acct_name = b'1234567890@developer.gserviceaccount.com'
-        http_request.return_value = (
-            httplib2.Response({'status': http_client.OK}), acct_name)
-        result = _get_service_account_email(http_request)
-        self.assertEqual(result, (None, acct_name.decode('utf-8')))
-        http_request.assert_called_once_with(
-            _DEFAULT_EMAIL_METADATA,
-            headers={'Metadata-Flavor': 'Google'})
-
-    @mock.patch.object(httplib2.Http, 'request')
-    def test_success_default_http(self, http_request):
-        # Don't make _from_bytes() work too hard.
-        acct_name = u'1234567890@developer.gserviceaccount.com'
-        http_request.return_value = (
-            httplib2.Response({'status': http_client.OK}), acct_name)
-        result = _get_service_account_email()
-        self.assertEqual(result, (None, acct_name))
-        http_request.assert_called_once_with(
-            _DEFAULT_EMAIL_METADATA,
-            headers={'Metadata-Flavor': 'Google'})
-
-    def test_failure(self):
-        http_request = mock.MagicMock()
-        response = httplib2.Response({'status': http_client.NOT_FOUND})
-        content = b'Not found'
-        http_request.return_value = (response, content)
-        result = _get_service_account_email(http_request)
-
-        self.assertEqual(result, (response, content))
-        http_request.assert_called_once_with(
-            _DEFAULT_EMAIL_METADATA,
-            headers={'Metadata-Flavor': 'Google'})
-
-
 if __name__ == '__main__':  # pragma: NO COVER
     unittest2.main()

tests/contrib/test_metadata.py

@@ -0,0 +1,97 @@
+# Copyright 2016 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.
+
+import datetime
+import httplib2
+import json
+import mock
+import unittest2
+
+from six.moves import http_client
+
+from oauth2client.contrib import _metadata
+
+PATH = 'instance/service-accounts/default'
+DATA = {'foo': 'bar'}
+EXPECTED_URL = (
+    'http://metadata.google.internal/computeMetadata/v1/instance'
+    '/service-accounts/default')
+EXPECTED_KWARGS = dict(headers=_metadata.METADATA_HEADERS)
+
+
+def request_mock(status, content_type, content):
+    return mock.MagicMock(return_value=(
+        httplib2.Response(
+            {'status': status, 'content-type': content_type}
+        ),
+        content.encode('utf-8')
+    ))
+
+
+class TestMetadata(unittest2.TestCase):
+
+    def test_get_success_json(self):
+        http_request = request_mock(
+            http_client.OK, 'application/json', json.dumps(DATA))
+        self.assertEqual(
+            _metadata.get(PATH, http_request=http_request),
+            DATA
+        )
+        http_request.assert_called_once_with(EXPECTED_URL, **EXPECTED_KWARGS)
+
+    def test_get_success_string(self):
+        http_request = request_mock(
+            http_client.OK, 'text/html', '<p>Hello World!</p>')
+        self.assertEqual(
+            _metadata.get(PATH, http_request=http_request),
+            '<p>Hello World!</p>'
+        )
+        http_request.assert_called_once_with(EXPECTED_URL, **EXPECTED_KWARGS)
+
+    def test_get_failure(self):
+        http_request = request_mock(
+            http_client.NOT_FOUND, 'text/html', '<p>Error</p>')
+        with self.assertRaises(httplib2.HttpLib2Error):
+            _metadata.get(PATH, http_request=http_request)
+
+        http_request.assert_called_once_with(EXPECTED_URL, **EXPECTED_KWARGS)
+
+    @mock.patch(
+        'oauth2client.contrib._metadata._UTCNOW',
+        return_value=datetime.datetime.min)
+    def test_get_token_success(self, now):
+        http_request = request_mock(
+            http_client.OK,
+            'application/json',
+            json.dumps({'access_token': 'a', 'expires_in': 100})
+        )
+        token, expiry = _metadata.get_token(http_request=http_request)
+        self.assertEqual(token, 'a')
+        self.assertEqual(
+            expiry, datetime.datetime.min + datetime.timedelta(seconds=100))
+        http_request.assert_called_once_with(
+            EXPECTED_URL+'/token',
+            **EXPECTED_KWARGS
+        )
+        now.assert_called_once_with()
+
+    def test_service_account_info(self):
+        http_request = request_mock(
+            http_client.OK, 'application/json', json.dumps(DATA))
+        info = _metadata.get_service_account_info(http_request=http_request)
+        self.assertEqual(info, DATA)
+        http_request.assert_called_once_with(
+            EXPECTED_URL+'?recursive=True',
+            **EXPECTED_KWARGS
+        )