Add compute engine metadata client

google/auth/_helpers.py

@@ -19,6 +19,7 @@
 import datetime
 
 import six
+from six.moves import urllib
 
 
 def utcnow():
@@ -88,3 +89,30 @@ def from_bytes(value):
     else:
         raise ValueError(
             '{0!r} could not be converted to unicode'.format(value))
+
+
+def update_query(url, params):
+    """Updates a URL's query parameters
+    Replaces any current values if they are already present in the URL.
+    Args:
+        url (str): The URL to update.
+        params (Mapping): A mapping of query parameter keys to values.
+    Returns:
+        str: The URL with updated query parameters.
+    """
+    # Split the URL into parts.
+    parts = urllib.parse.urlparse(url)
+    # Parse the query string.
+    query_params = urllib.parse.parse_qs(parts.query)
+    # Update the query parameters with the new parameters.
+    query_params.update(params)
+    # Remove any None values.
+    query_params = {
+        key: value for key, value
+        in six.iteritems(query_params)
+        if value is not None}
+    # Re-encoded the query string.
+    new_query = urllib.parse.urlencode(query_params, doseq=True)
+    # Unsplit the url.
+    new_parts = parts[:4] + (new_query,) + parts[5:]
+    return urllib.parse.urlunparse(new_parts)

google/auth/compute_engine/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2016 Google Inc.
+#
+# 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.

google/auth/compute_engine/_metadata.py

@@ -0,0 +1,169 @@
+# Copyright 2016 Google Inc.
+#
+# 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 for more details.
+"""
+
+import datetime
+import json
+import os
+
+from six.moves import http_client
+from six.moves.urllib import parse as urlparse
+
+from google.auth import _helpers
+from google.auth import exceptions
+
+_METADATA_ROOT = 'http://metadata.google.internal/computeMetadata/v1/'
+
+# This is used to ping the metadata server, it avoids the cost of a DNS
+# lookup.
+_METADATA_IP_ROOT = 'http://169.254.169.254'
+_METADATA_FLAVOR_HEADER = 'metdata-flavor'
+_METADATA_FLAVOR_VALUE = 'Google'
+_METADATA_HEADERS = {_METADATA_FLAVOR_HEADER: _METADATA_FLAVOR_VALUE}
+
+# Timeout in seconds to wait for the GCE metadata server when detecting the
+# GCE environment.
+try:
+    _METADATA_DEFAULT_TIMEOUT = int(os.getenv('GCE_METADATA_TIMEOUT', 3))
+except ValueError:  # pragma: NO COVER
+    _METADATA_DEFAULT_TIMEOUT = 3
+
+
+def ping(request, timeout=_METADATA_DEFAULT_TIMEOUT):
+    """Checks to see if the metadata server is available.
+
+    Args:
+        request (google.auth.transport.Request): A callable used to make
+            HTTP requests.
+        timeout (int): How long to wait for the metadata server to respond.
+
+    Returns:
+        bool: True if the metadata server is reachable, False otherwise.
+    """
+    # NOTE: The explicit ``timeout`` is a workaround. The underlying
+    #       issue is that resolving an unknown host on some networks will take
+    #       20-30 seconds; making this timeout short fixes the issue, but
+    #       could lead to false negatives in the event that we are on GCE, but
+    #       the metadata resolution was particularly slow. The latter case is
+    #       "unlikely".
+    try:
+        response = request(
+            url=_METADATA_IP_ROOT, method='GET', headers=_METADATA_HEADERS,
+            timeout=timeout)
+
+        metadata_flavor = response.headers.get(_METADATA_FLAVOR_HEADER)
+        return (response.status == http_client.OK and
+                metadata_flavor == _METADATA_FLAVOR_VALUE)
+
+    except exceptions.TransportError:
+        # logger.info('Timeout attempting to reach GCE metadata service.')
+        return False
+
+
+def get(request, path, root=_METADATA_ROOT, recursive=None):
+    """Fetch a resource from the metadata server.
+
+    Args:
+        request (google.auth.transport.Request): A callable used to make
+            HTTP requests.
+        path (str): The resource to retrieve. For example,
+            ``'instance/service-accounts/defualt'``.
+        root (str): The full path to the metadata server root.
+        recursive (bool): Whether to do a recursive query of metadata. See
+            https://cloud.google.com/compute/docs/metadata#aggcontents for more
+            details.
+
+    Returns:
+        Union[Mapping, str]: If the metadata server returns JSON, a mapping of
+            the decoded JSON is return. Otherwise, the response content is
+            returned as a string.
+
+    Raises:
+        google.auth.exceptions.TransportError: if an error occurred while
+            retrieving metadata.
+    """
+    url = urlparse.urljoin(root, path)
+    url = _helpers.update_query(url, {'recursive': recursive})
+
+    response = request(url=url, method='GET', headers=_METADATA_HEADERS)
+
+    if response.status == http_client.OK:
+        content = _helpers.from_bytes(response.data)
+        if response.headers['content-type'] == 'application/json':
+            return json.loads(content)
+        else:
+            return content
+    else:
+        raise exceptions.TransportError(
+            'Failed to retrieve {} from the Google Compute Engine'
+            'metadata service. Status: {} Response:\n{}'.format(
+                url, response.status, response.data), response)
+
+
+def get_service_account_info(request, service_account='default'):
+    """Get information about a service account from the metadata server.
+
+    Args:
+        request (google.auth.transport.Request): A callable used to make
+            HTTP requests.
+        service_account (str): The string 'default' or a service account email
+            address. The determines which service account for which to acquire
+            information.
+
+    Returns:
+        Mapping: The service account's information, for example::
+
+            {
+                'email': '...',
+                'scopes': ['scope', ...],
+                'aliases': ['default', '...']
+            }
+
+    Raises:
+        google.auth.exceptions.TransportError: if an error occurred while
+            retrieving metadata.
+    """
+    return get(
+        request,
+        'instance/service-accounts/{0}/'.format(service_account),
+        recursive=True)
+
+
+def get_service_account_token(request, service_account='default'):
+    """Get the OAuth 2.0 access token for a service account.
+
+    Args:
+        request (google.auth.transport.Request): A callable used to make
+            HTTP requests.
+        service_account (str): The string 'default' or a service account email
+            address. The determines which service account for which to acquire
+            an access token.
+
+    Returns:
+        Union[str, datetime]: The access token and its expiration.
+
+    Raises:
+        google.auth.exceptions.TransportError: if an error occurred while
+            retrieving metadata.
+    """
+    token_json = get(
+        request,
+        'instance/service-accounts/{0}/token'.format(service_account))
+    token_expiry = _helpers.utcnow() + datetime.timedelta(
+        seconds=token_json['expires_in'])
+    return token_json['access_token'], token_expiry

tests/compute_engine/test__metadata.py

@@ -0,0 +1,140 @@
+# Copyright 2016 Google Inc.
+#
+# 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 json
+
+import mock
+import pytest
+from six.moves import http_client
+
+from google.auth import _helpers
+from google.auth import exceptions
+from google.auth.compute_engine import _metadata
+
+PATH = 'instance/service-accounts/default'
+
+
+@pytest.fixture
+def mock_request():
+    request_mock = mock.Mock()
+
+    def set_response(data, status=http_client.OK, headers=None):
+        response = mock.Mock()
+        response.status = status
+        response.data = _helpers.to_bytes(data)
+        response.headers = headers or {}
+        request_mock.return_value = response
+        return request_mock
+
+    yield set_response
+
+
+def test_ping_success(mock_request):
+    request_mock = mock_request('', headers=_metadata._METADATA_HEADERS)
+
+    assert _metadata.ping(request_mock)
+
+    request_mock.assert_called_once_with(
+        method='GET',
+        url=_metadata._METADATA_IP_ROOT,
+        headers=_metadata._METADATA_HEADERS,
+        timeout=_metadata._METADATA_DEFAULT_TIMEOUT)
+
+
+def test_ping_failure_bad_flavor(mock_request):
+    request_mock = mock_request(
+        '', headers={_metadata._METADATA_FLAVOR_HEADER: 'meep'})
+
+    assert not _metadata.ping(request_mock)
+
+
+def test_ping_failure_connection_failed(mock_request):
+    request_mock = mock_request('')
+    request_mock.side_effect = exceptions.TransportError()
+
+    assert not _metadata.ping(request_mock)
+
+
+def test_get_success_json(mock_request):
+    data = json.dumps({'foo': 'bar'})
+    request_mock = mock_request(
+        data, headers={'content-type': 'application/json'})
+
+    result = _metadata.get(request_mock, PATH)
+
+    request_mock.assert_called_once_with(
+        method='GET',
+        url=_metadata._METADATA_ROOT + PATH,
+        headers=_metadata._METADATA_HEADERS)
+    assert result['foo'] == 'bar'
+
+
+def test_get_success_text(mock_request):
+    data = 'foobar'
+    request_mock = mock_request(data, headers={'content-type': 'text/plain'})
+
+    result = _metadata.get(request_mock, PATH)
+
+    request_mock.assert_called_once_with(
+        method='GET',
+        url=_metadata._METADATA_ROOT + PATH,
+        headers=_metadata._METADATA_HEADERS)
+    assert result == data
+
+
+def test_get_failure(mock_request):
+    request_mock = mock_request(
+        'Metadata error', status=http_client.NOT_FOUND)
+
+    with pytest.raises(exceptions.TransportError) as excinfo:
+        _metadata.get(request_mock, PATH)
+
+    assert excinfo.match(r'Metadata error')
+
+    request_mock.assert_called_once_with(
+        method='GET',
+        url=_metadata._METADATA_ROOT + PATH,
+        headers=_metadata._METADATA_HEADERS)
+
+
+@mock.patch('google.auth._helpers.utcnow', return_value=datetime.datetime.min)
+def test_get_service_account_token(utcnow, mock_request):
+    request_mock = mock_request(
+        json.dumps({'access_token': 'token', 'expires_in': 500}),
+        headers={'content-type': 'application/json'})
+
+    token, expiry = _metadata.get_service_account_token(request_mock)
+
+    request_mock.assert_called_once_with(
+        method='GET',
+        url=_metadata._METADATA_ROOT + PATH + '/token',
+        headers=_metadata._METADATA_HEADERS)
+    assert token == 'token'
+    assert expiry == utcnow() + datetime.timedelta(seconds=500)
+
+
+def test_get_service_account_info(mock_request):
+    request_mock = mock_request(
+        json.dumps({'foo': 'bar'}),
+        headers={'content-type': 'application/json'})
+
+    info = _metadata.get_service_account_info(request_mock)
+
+    request_mock.assert_called_once_with(
+        method='GET',
+        url=_metadata._METADATA_ROOT + PATH + '/?recursive=True',
+        headers=_metadata._METADATA_HEADERS)
+
+    assert info['foo'] == 'bar'