src/tools/telemetry/telemetry/util/cloud_storage.py

   1 # Copyright 2014 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Wrappers for gsutil, for basic interaction with Google Cloud Storage."""
   6
   7 import contextlib
   8 import cStringIO
   9 import hashlib
  10 import logging
  11 import os
  12 import subprocess
  13 import sys
  14 import tarfile
  15 import urllib2
  16
  17 from telemetry.core import platform
  18 from telemetry.util import path
  19
  20
  21 PUBLIC_BUCKET = 'chromium-telemetry'
  22 PARTNER_BUCKET = 'chrome-partner-telemetry'
  23 INTERNAL_BUCKET = 'chrome-telemetry'
  24
  25
  26 BUCKET_ALIASES = {
  27     'public': PUBLIC_BUCKET,
  28     'partner': PARTNER_BUCKET,
  29     'internal': INTERNAL_BUCKET,
  30 }
  31
  32
  33 _GSUTIL_URL = 'http://storage.googleapis.com/pub/gsutil.tar.gz'
  34 _DOWNLOAD_PATH = os.path.join(path.GetTelemetryDir(), 'third_party', 'gsutil')
  35 # TODO(tbarzic): A workaround for http://crbug.com/386416 and
  36 #     http://crbug.com/359293. See |_RunCommand|.
  37 _CROS_GSUTIL_HOME_WAR = '/home/chromeos-test/'
  38
  39
  40 class CloudStorageError(Exception):
  41   @staticmethod
  42   def _GetConfigInstructions(gsutil_path):
  43     if SupportsProdaccess(gsutil_path) and _FindExecutableInPath('prodaccess'):
  44       return 'Run prodaccess to authenticate.'
  45     else:
  46       if platform.GetHostPlatform().GetOSName() == 'chromeos':
  47         gsutil_path = ('HOME=%s %s' % (_CROS_GSUTIL_HOME_WAR, gsutil_path))
  48       return ('To configure your credentials:\n'
  49               '  1. Run "%s config" and follow its instructions.\n'
  50               '  2. If you have a @google.com account, use that account.\n'
  51               '  3. For the project-id, just enter 0.' % gsutil_path)
  52
  53
  54 class PermissionError(CloudStorageError):
  55   def __init__(self, gsutil_path):
  56     super(PermissionError, self).__init__(
  57         'Attempted to access a file from Cloud Storage but you don\'t '
  58         'have permission. ' + self._GetConfigInstructions(gsutil_path))
  59
  60
  61 class CredentialsError(CloudStorageError):
  62   def __init__(self, gsutil_path):
  63     super(CredentialsError, self).__init__(
  64         'Attempted to access a file from Cloud Storage but you have no '
  65         'configured credentials. ' + self._GetConfigInstructions(gsutil_path))
  66
  67
  68 class NotFoundError(CloudStorageError):
  69   pass
  70
  71
  72 class ServerError(CloudStorageError):
  73   pass
  74
  75
  76 # TODO(tonyg/dtu): Can this be replaced with distutils.spawn.find_executable()?
  77 def _FindExecutableInPath(relative_executable_path, *extra_search_paths):
  78   search_paths = list(extra_search_paths) + os.environ['PATH'].split(os.pathsep)
  79   for search_path in search_paths:
  80     executable_path = os.path.join(search_path, relative_executable_path)
  81     if path.IsExecutable(executable_path):
  82       return executable_path
  83   return None
  84
  85
  86 def _DownloadGsutil():
  87   logging.info('Downloading gsutil')
  88   with contextlib.closing(urllib2.urlopen(_GSUTIL_URL, timeout=60)) as response:
  89     with tarfile.open(fileobj=cStringIO.StringIO(response.read())) as tar_file:
  90       tar_file.extractall(os.path.dirname(_DOWNLOAD_PATH))
  91   logging.info('Downloaded gsutil to %s' % _DOWNLOAD_PATH)
  92
  93   return os.path.join(_DOWNLOAD_PATH, 'gsutil')
  94
  95
  96 def FindGsutil():
  97   """Return the gsutil executable path. If we can't find it, download it."""
  98   # Look for a depot_tools installation.
  99   # FIXME: gsutil in depot_tools is not working correctly. crbug.com/413414
 100   #gsutil_path = _FindExecutableInPath(
 101   #    os.path.join('third_party', 'gsutil', 'gsutil'), _DOWNLOAD_PATH)
 102   #if gsutil_path:
 103   #  return gsutil_path
 104
 105   # Look for a gsutil installation.
 106   gsutil_path = _FindExecutableInPath('gsutil', _DOWNLOAD_PATH)
 107   if gsutil_path:
 108     return gsutil_path
 109
 110   # Failed to find it. Download it!
 111   return _DownloadGsutil()
 112
 113
 114 def SupportsProdaccess(gsutil_path):
 115   with open(gsutil_path, 'r') as gsutil:
 116     return 'prodaccess' in gsutil.read()
 117
 118
 119 def _RunCommand(args):
 120   gsutil_path = FindGsutil()
 121
 122   # On cros device, as telemetry is running as root, home will be set to /root/,
 123   # which is not writable. gsutil will attempt to create a download tracker dir
 124   # in home dir and fail. To avoid this, override HOME dir to something writable
 125   # when running on cros device.
 126   #
 127   # TODO(tbarzic): Figure out a better way to handle gsutil on cros.
 128   #     http://crbug.com/386416, http://crbug.com/359293.
 129   gsutil_env = None
 130   if platform.GetHostPlatform().GetOSName() == 'chromeos':
 131     gsutil_env = os.environ.copy()
 132     gsutil_env['HOME'] = _CROS_GSUTIL_HOME_WAR
 133
 134   gsutil = subprocess.Popen([sys.executable, gsutil_path] + args,
 135                             stdout=subprocess.PIPE, stderr=subprocess.PIPE,
 136                             env=gsutil_env)
 137   stdout, stderr = gsutil.communicate()
 138
 139   if gsutil.returncode:
 140     if stderr.startswith((
 141         'You are attempting to access protected data with no configured',
 142         'Failure: No handler was ready to authenticate.')):
 143       raise CredentialsError(gsutil_path)
 144     if 'status=403' in stderr or 'status 403' in stderr:
 145       raise PermissionError(gsutil_path)
 146     if (stderr.startswith('InvalidUriError') or 'No such object' in stderr or
 147         'No URLs matched' in stderr):
 148       raise NotFoundError(stderr)
 149     if '500 Internal Server Error' in stderr:
 150       raise ServerError(stderr)
 151     raise CloudStorageError(stderr)
 152
 153   return stdout
 154
 155
 156 def List(bucket):
 157   query = 'gs://%s/' % bucket
 158   stdout = _RunCommand(['ls', query])
 159   return [url[len(query):] for url in stdout.splitlines()]
 160
 161
 162 def Exists(bucket, remote_path):
 163   try:
 164     _RunCommand(['ls', 'gs://%s/%s' % (bucket, remote_path)])
 165     return True
 166   except NotFoundError:
 167     return False
 168
 169
 170 def Move(bucket1, bucket2, remote_path):
 171   url1 = 'gs://%s/%s' % (bucket1, remote_path)
 172   url2 = 'gs://%s/%s' % (bucket2, remote_path)
 173   logging.info('Moving %s to %s' % (url1, url2))
 174   _RunCommand(['mv', url1, url2])
 175
 176
 177 def Delete(bucket, remote_path):
 178   url = 'gs://%s/%s' % (bucket, remote_path)
 179   logging.info('Deleting %s' % url)
 180   _RunCommand(['rm', url])
 181
 182
 183 def Get(bucket, remote_path, local_path):
 184   url = 'gs://%s/%s' % (bucket, remote_path)
 185   logging.info('Downloading %s to %s' % (url, local_path))
 186   try:
 187     _RunCommand(['cp', url, local_path])
 188   except ServerError:
 189     logging.info('Cloud Storage server error, retrying download')
 190     _RunCommand(['cp', url, local_path])
 191
 192
 193 def Insert(bucket, remote_path, local_path, publicly_readable=False):
 194   url = 'gs://%s/%s' % (bucket, remote_path)
 195   command_and_args = ['cp']
 196   extra_info = ''
 197   if publicly_readable:
 198     command_and_args += ['-a', 'public-read']
 199     extra_info = ' (publicly readable)'
 200   command_and_args += [local_path, url]
 201   logging.info('Uploading %s to %s%s' % (local_path, url, extra_info))
 202   _RunCommand(command_and_args)
 203
 204
 205 def GetIfChanged(file_path, bucket=None):
 206   """Gets the file at file_path if it has a hash file that doesn't match.
 207
 208   If the file is not in Cloud Storage, log a warning instead of raising an
 209   exception. We assume that the user just hasn't uploaded the file yet.
 210
 211   Returns:
 212     True if the binary was changed.
 213   """
 214   hash_path = file_path + '.sha1'
 215   if not os.path.exists(hash_path):
 216     logging.warning('Hash file not found: %s' % hash_path)
 217     return False
 218
 219   expected_hash = ReadHash(hash_path)
 220   if os.path.exists(file_path) and CalculateHash(file_path) == expected_hash:
 221     return False
 222
 223   if bucket:
 224     buckets = [bucket]
 225   else:
 226     buckets = [PUBLIC_BUCKET, PARTNER_BUCKET, INTERNAL_BUCKET]
 227
 228   for bucket in buckets:
 229     try:
 230       Get(bucket, expected_hash, file_path)
 231       return True
 232     except NotFoundError:
 233       continue
 234
 235   logging.warning('Unable to find file in Cloud Storage: %s', file_path)
 236   return False
 237
 238
 239 def CalculateHash(file_path):
 240   """Calculates and returns the hash of the file at file_path."""
 241   sha1 = hashlib.sha1()
 242   with open(file_path, 'rb') as f:
 243     while True:
 244       # Read in 1mb chunks, so it doesn't all have to be loaded into memory.
 245       chunk = f.read(1024*1024)
 246       if not chunk:
 247         break
 248       sha1.update(chunk)
 249   return sha1.hexdigest()
 250
 251
 252 def ReadHash(hash_path):
 253   with open(hash_path, 'rb') as f:
 254     return f.read(1024).rstrip()