Index: scripts/slave/gce.py |
diff --git a/scripts/slave/gce.py b/scripts/slave/gce.py |
new file mode 100755 |
index 0000000000000000000000000000000000000000..398c8a1ece5a285a26a911955e2a458cb38adc83 |
--- /dev/null |
+++ b/scripts/slave/gce.py |
@@ -0,0 +1,88 @@ |
+# Copyright 2015 The Chromium OS Authors. All rights reserved. |
+# Use of this source code is governed by a BSD-style license that can be |
+# found in the LICENSE file. |
+ |
+""" |
+Utilities for interfacing with Google Compute Engine. |
+""" |
+ |
+import httplib |
+import json |
+import logging |
+import socket |
+import time |
+import urlparse |
+ |
+ |
+LOGGER = logging.getLogger('gce') |
+TRY_LIMIT = 5 |
+ |
+ |
+class Authenticator(object): |
+ """Authenticator implementation that uses GCE metadata service for token. |
+ """ |
+ |
+ _INFO_URL = 'http://metadata.google.internal' |
+ _ACQUIRE_URL = ('http://metadata/computeMetadata/v1/instance/' |
+ 'service-accounts/default/token') |
+ _ACQUIRE_HEADERS = {"Metadata-Flavor": "Google"} |
+ |
+ _cache_is_gce = None |
+ _token_cache = None |
+ _token_expiration = None |
+ |
+ @classmethod |
+ def is_gce(cls): |
+ if cls._cache_is_gce is None: |
+ cls._cache_is_gce = cls._test_is_gce() |
+ return cls._cache_is_gce |
+ |
+ @classmethod |
+ def _test_is_gce(cls): |
+ # Based on https://cloud.google.com/compute/docs/metadata#runninggce |
+ try: |
+ resp = cls._get(cls._INFO_URL) |
+ except socket.error: |
+ # Could not resolve URL. |
+ return False |
+ return resp.getheader('Metadata-Flavor', None) == 'Google' |
+ |
+ @staticmethod |
+ def _get(url, **kwargs): |
+ next_delay_sec = 1 |
+ for i in xrange(TRY_LIMIT): |
+ if i > 0: |
+ # Retry server error status codes. |
+ LOGGER.info('Encountered server error; retrying after %d second(s).', |
+ next_delay_sec) |
+ time.sleep(next_delay_sec) |
+ next_delay_sec *= 2 |
+ |
+ p = urlparse.urlparse(url) |
+ c = GetConnectionClass(protocol=p.scheme)(p.netloc) |
+ c.request('GET', url, **kwargs) |
+ resp = c.getresponse() |
+ LOGGER.debug('GET [%s] #%d/%d (%d)', url, i+1, TRY_LIMIT, resp.status) |
+ if resp.status < httplib.INTERNAL_SERVER_ERROR: |
+ return resp |
+ |
+ |
+ @classmethod |
+ def _get_token_dict(cls): |
+ if cls._token_cache: |
+ # If it expires within 25 seconds, refresh. |
+ if cls._token_expiration < time.time() - 25: |
+ return cls._token_cache |
+ |
+ resp = cls._get(cls._ACQUIRE_URL, headers=cls._ACQUIRE_HEADERS) |
+ if resp.status != httplib.OK: |
+ return None |
+ cls._token_cache = json.load(resp) |
+ cls._token_expiration = cls._token_cache['expires_in'] + time.time() |
+ return cls._token_cache |
+ |
+ def get_auth_header(self, _host): |
+ token_dict = self._get_token_dict() |
+ if not token_dict: |
+ return None |
+ return '%(token_type)s %(access_token)s' % token_dict |