Add compute engine metadata client (#11)
diff --git a/google/auth/_helpers.py b/google/auth/_helpers.py
index 0a62209..d4bf8f1 100644
--- a/google/auth/_helpers.py
+++ b/google/auth/_helpers.py
@@ -19,6 +19,7 @@
import datetime
import six
+from six.moves import urllib
def utcnow():
@@ -88,3 +89,49 @@
else:
raise ValueError(
'{0!r} could not be converted to unicode'.format(value))
+
+
+def update_query(url, params, remove=None):
+ """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[str, str]): A mapping of query parameter
+ keys to values.
+ remove (Sequence[str]): Parameters to remove from the query string.
+
+ Returns:
+ str: The URL with updated query parameters.
+
+ Examples:
+
+ >>> url = 'http://example.com?a=1'
+ >>> update_query(url, {'a': '2'})
+ http://example.com?a=2
+ >>> update_query(url, {'b': '3'})
+ http://example.com?a=1&b=3
+ >> update_query(url, {'b': '3'}, remove=['a'])
+ http://example.com?b=3
+
+ """
+ if remove is None:
+ remove = []
+
+ # 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 values specified in remove.
+ query_params = {
+ key: value for key, value
+ in six.iteritems(query_params)
+ if key not in remove}
+ # Re-encoded the query string.
+ new_query = urllib.parse.urlencode(query_params, doseq=True)
+ # Unsplit the url.
+ new_parts = parts._replace(query=new_query)
+ return urllib.parse.urlunparse(new_parts)
diff --git a/google/auth/compute_engine/__init__.py b/google/auth/compute_engine/__init__.py
new file mode 100644
index 0000000..e3a7f6c
--- /dev/null
+++ b/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.
diff --git a/google/auth/compute_engine/_metadata.py b/google/auth/compute_engine/_metadata.py
new file mode 100644
index 0000000..28d57f6
--- /dev/null
+++ b/google/auth/compute_engine/_metadata.py
@@ -0,0 +1,178 @@
+# 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:
+ return False
+
+
+def get(request, path, root=_METADATA_ROOT, recursive=False):
+ """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.
+ """
+ base_url = urlparse.urljoin(root, path)
+ query_params = {}
+
+ if recursive:
+ query_params['recursive'] = 'true'
+
+ url = _helpers.update_query(base_url, query_params)
+
+ 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':
+ try:
+ return json.loads(content)
+ except ValueError:
+ raise exceptions.TransportError(
+ 'Received invalid JSON from the Google Compute Engine'
+ 'metadata service: {:.20}'.format(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
diff --git a/pylintrc.tests b/pylintrc.tests
index 73b2766..de1964f 100644
--- a/pylintrc.tests
+++ b/pylintrc.tests
@@ -125,7 +125,7 @@
# DEFAULT: good-names=i,j,k,ex,Run,_
# RATIONALE: 'fh' is a well-known file handle variable name.
good-names = i, j, k, ex, Run, _,
- fh,
+ fh
# Regular expression matching correct method names
# DEFAULT: method-rgx=[a-z_][a-z0-9_]{2,30}$
diff --git a/tests/compute_engine/__init__.py b/tests/compute_engine/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/tests/compute_engine/__init__.py
diff --git a/tests/compute_engine/test__metadata.py b/tests/compute_engine/test__metadata.py
new file mode 100644
index 0000000..4cc3e55
--- /dev/null
+++ b/tests/compute_engine/test__metadata.py
@@ -0,0 +1,159 @@
+# 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):
+ key, value = 'foo', 'bar'
+
+ data = json.dumps({key: value})
+ 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[key] == value
+
+
+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)
+
+
+def test_get_failure_bad_json(mock_request):
+ request_mock = mock_request(
+ '{', headers={'content-type': 'application/json'})
+
+ with pytest.raises(exceptions.TransportError) as excinfo:
+ _metadata.get(request_mock, PATH)
+
+ assert excinfo.match(r'invalid JSON')
+
+ 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):
+ ttl = 500
+ request_mock = mock_request(
+ json.dumps({'access_token': 'token', 'expires_in': ttl}),
+ 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=ttl)
+
+
+def test_get_service_account_info(mock_request):
+ key, value = 'foo', 'bar'
+ request_mock = mock_request(
+ json.dumps({key: value}),
+ 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[key] == value
diff --git a/tests/test__helpers.py b/tests/test__helpers.py
index c2bc4a7..d475fc4 100644
--- a/tests/test__helpers.py
+++ b/tests/test__helpers.py
@@ -15,6 +15,7 @@
import datetime
import pytest
+from six.moves import urllib
from google.auth import _helpers
@@ -60,3 +61,35 @@
def test_from_bytes_with_nonstring_type():
with pytest.raises(ValueError):
_helpers.from_bytes(object())
+
+
+def _assert_query(url, expected):
+ parts = urllib.parse.urlsplit(url)
+ query = urllib.parse.parse_qs(parts.query)
+ assert query == expected
+
+
+def test_update_query_params_no_params():
+ uri = 'http://www.google.com'
+ updated = _helpers.update_query(uri, {'a': 'b'})
+ assert updated == uri + '?a=b'
+
+
+def test_update_query_existing_params():
+ uri = 'http://www.google.com?x=y'
+ updated = _helpers.update_query(uri, {'a': 'b', 'c': 'd&'})
+ _assert_query(updated, {'x': ['y'], 'a': ['b'], 'c': ['d&']})
+
+
+def test_update_query_replace_param():
+ base_uri = 'http://www.google.com'
+ uri = base_uri + '?x=a'
+ updated = _helpers.update_query(uri, {'x': 'b', 'y': 'c'})
+ _assert_query(updated, {'x': ['b'], 'y': ['c']})
+
+
+def test_update_query_remove_param():
+ base_uri = 'http://www.google.com'
+ uri = base_uri + '?x=a'
+ updated = _helpers.update_query(uri, {'y': 'c'}, remove=['x'])
+ _assert_query(updated, {'y': ['c']})