1 files changed, 115 insertions, 0 deletions
diff --git a/lib/cipd.py b/lib/cipd.py
new file mode 100644
index 000000000..427e02755
--- /dev/null
+++ b/lib/cipd.py
@@ -0,0 +1,115 @@
+# 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.
+
+"""Module to download and run the CIPD client.
+
+CIPD is the Chrome Infra Package Deployer, a simple method of resolving
+package/architecture/version tuples into GStorage links and installing them.
+
+"""
+from __future__ import print_function
+
+import hashlib
+import json
+import os
+import pprint
+import urllib
+import urlparse
+
+import third_party.httplib2 as httplib2
+
+import chromite.lib.cros_logging as log
+from chromite.lib import cache
+from chromite.lib import osutils
+from chromite.lib import path_util
+
+
+# The version of CIPD to download.
+# TODO(phobbs) we could make a call to the 'resolveVersion' endpoint
+#   to resolve 'latest' into an instance_id for us.
+CIPD_INSTANCE_ID = '03f354ad7a6031c7924d9b69a85f83269cc3c2e0'
+CIPD_PACKAGE = 'infra/tools/cipd/linux-amd64'
+
+CHROME_INFRA_PACKAGES_API_BASE = (
+    'https://chrome-infra-packages.appspot.com/_ah/api/repo/v1/')
+
+
+def _ChromeInfraRequest(endpoint, request_args=None):
+  """Makes a request to the Chrome Infra Packages API with httplib2.
+
+  Args:
+    endpoint: The endpoint to make a request to.
+    request_args: Keyword arguments to put into the request string.
+
+  Returns:
+    A tuple of (headers, content) returned by the server. The body content is
+    assumed to be JSON.
+  """
+  uri = ''.join([CHROME_INFRA_PACKAGES_API_BASE,
+                 endpoint,
+                 '?',
+                 urllib.urlencode(request_args or {})])
+  result = httplib2.Http().request(uri=uri)
+  try:
+    return result[0], json.loads(result[1])
+  except Exception as e:
+    e.message = 'Encountered exception requesting "%s":\n' + e.message
+    raise
+
+
+def _DownloadCIPD(instance_id):
+  """Finds the CIPD download link and requests the binary.
+
+  The 'client' endpoit of the chrome infra packages API responds with a sha1 and
+  a Google Storage link. After downloading the binary, we validate that the sha1
+  of the response and return it.
+
+  Args:
+    instance_id: The version of CIPD to download.
+
+  Returns:
+    the CIPD binary as a string.
+  """
+  args = {'instance_id': instance_id, 'package_name': CIPD_PACKAGE}
+  _, body = _ChromeInfraRequest('client', request_args=args)
+  if not 'client_binary' in body:
+    log.error(
+        'Error requesting the link to download CIPD from. Got:\n%s',
+        pprint.pformat(body))
+
+  http = httplib2.Http(cache='.cache')
+  response, binary = http.request(uri=body['client_binary']['fetch_url'])
+  assert response['status'] == '200', (
+      'Got a %s response from Google Storage.' % response['status'])
+  digest = unicode(hashlib.sha1(binary).hexdigest())
+  assert digest == body['client_binary']['sha1'], (
+      'The binary downloaded does not match the expected SHA1.')
+  return binary
+
+
+class CipdCache(cache.RemoteCache):
+  """Supports caching of the CIPD download."""
+  def _Fetch(self, key, path):
+    instance_id = urlparse.urlparse(key).netloc
+    binary = _DownloadCIPD(instance_id)
+    log.info('Fetched CIPD package %s:%s', CIPD_PACKAGE, instance_id)
+    osutils.WriteFile(path, binary)
+    os.chmod(path, 0755)
+
+
+def GetCIPDFromCache(instance_id=CIPD_INSTANCE_ID):
+  """Checks the cache, downloading CIPD if it is missing.
+
+  Args:
+    instance_id: The version of CIPD to download. Default CIPD_INSTANCE_ID
+
+  Returns:
+    Path to the CIPD binary.
+  """
+  cache_dir = os.path.join(path_util.GetCacheDir(), 'cipd')
+  bin_cache = CipdCache(cache_dir)
+  key = (instance_id,)
+  ref = bin_cache.Lookup(key)
+  ref.SetDefault('cipd://' + instance_id)
+  return ref.path