Add LogDog bootstrapping to `annotated_run.py`.

scripts/slave/gce.py

+# Copyright 2015 The Chromium 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
+
+
+def GetConnectionClass(protocol=None):
+  if protocol is None:
+    protocol = 'https'
+  if protocol == 'https':
+    return httplib.HTTPSConnection
+  elif protocol == 'http':
+    return httplib.HTTPConnection
+  else:
+    raise RuntimeError(
+        "Don't know how to work with protocol '%s'" % protocol)
+