tools/telemetry/catapult_base/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 collections
   8 import contextlib
   9 import cStringIO
  10 import hashlib
  11 import logging
  12 import os
  13 import subprocess
  14 import sys
  15 import tarfile
  16 import urllib2
  17
  18 from telemetry.core import util
  19 from telemetry import decorators
  20 from telemetry.internal.util import path
  21
  22
  23 PUBLIC_BUCKET = 'chromium-telemetry'
  24 PARTNER_BUCKET = 'chrome-partner-telemetry'
  25 INTERNAL_BUCKET = 'chrome-telemetry'
  26
  27
  28 # Uses ordered dict to make sure that bucket's key-value items are ordered from
  29 # the most open to the most restrictive.
  30 BUCKET_ALIASES = collections.OrderedDict((
  31     ('public', PUBLIC_BUCKET),
  32     ('partner', PARTNER_BUCKET),
  33     ('internal', INTERNAL_BUCKET),
  34 ))
  35
  36
  37 _GSUTIL_URL = 'http://storage.googleapis.com/pub/gsutil.tar.gz'
  38 _DOWNLOAD_PATH = os.path.join(path.GetTelemetryDir(), 'third_party', 'gsutil')
  39 # TODO(tbarzic): A workaround for http://crbug.com/386416 and
  40 #     http://crbug.com/359293. See |_RunCommand|.
  41 _CROS_GSUTIL_HOME_WAR = '/home/chromeos-test/'
  42
  43
  44 class CloudStorageError(Exception):
  45   @staticmethod
  46   def _GetConfigInstructions(gsutil_path):
  47     if SupportsProdaccess(gsutil_path) and _FindExecutableInPath('prodaccess'):
  48       return 'Run prodaccess to authenticate.'
  49     else:
  50       if util.IsRunningOnCrosDevice():
  51         gsutil_path = ('HOME=%s %s' % (_CROS_GSUTIL_HOME_WAR, gsutil_path))
  52       return ('To configure your credentials:\n'
  53               '  1. Run "%s config" and follow its instructions.\n'
  54               '  2. If you have a @google.com account, use that account.\n'
  55               '  3. For the project-id, just enter 0.' % gsutil_path)
  56
  57
  58 class PermissionError(CloudStorageError):
  59   def __init__(self, gsutil_path):
  60     super(PermissionError, self).__init__(
  61         'Attempted to access a file from Cloud Storage but you don\'t '
  62         'have permission. ' + self._GetConfigInstructions(gsutil_path))
  63
  64
  65 class CredentialsError(CloudStorageError):
  66   def __init__(self, gsutil_path):
  67     super(CredentialsError, self).__init__(
  68         'Attempted to access a file from Cloud Storage but you have no '
  69         'configured credentials. ' + self._GetConfigInstructions(gsutil_path))
  70
  71
  72 class NotFoundError(CloudStorageError):
  73   pass
  74
  75
  76 class ServerError(CloudStorageError):
  77   pass
  78
  79
  80 # TODO(tonyg/dtu): Can this be replaced with distutils.spawn.find_executable()?
  81 def _FindExecutableInPath(relative_executable_path, *extra_search_paths):
  82   search_paths = list(extra_search_paths) + os.environ['PATH'].split(os.pathsep)
  83   for search_path in search_paths:
  84     executable_path = os.path.join(search_path, relative_executable_path)
  85     if path.IsExecutable(executable_path):
  86       return executable_path
  87   return None
  88
  89
  90 def _DownloadGsutil():
  91   logging.info('Downloading gsutil')
  92   with contextlib.closing(urllib2.urlopen(_GSUTIL_URL, timeout=60)) as response:
  93     with tarfile.open(fileobj=cStringIO.StringIO(response.read())) as tar_file:
  94       tar_file.extractall(os.path.dirname(_DOWNLOAD_PATH))
  95   logging.info('Downloaded gsutil to %s' % _DOWNLOAD_PATH)
  96
  97   return os.path.join(_DOWNLOAD_PATH, 'gsutil')
  98
  99
 100 def FindGsutil():
 101   """Return the gsutil executable path. If we can't find it, download it."""
 102   # Look for a depot_tools installation.
 103   # FIXME: gsutil in depot_tools is not working correctly. crbug.com/413414
 104   #gsutil_path = _FindExecutableInPath(
 105   #    os.path.join('third_party', 'gsutil', 'gsutil'), _DOWNLOAD_PATH)
 106   #if gsutil_path:
 107   #  return gsutil_path
 108
 109   # Look for a gsutil installation.
 110   gsutil_path = _FindExecutableInPath('gsutil', _DOWNLOAD_PATH)
 111   if gsutil_path:
 112     return gsutil_path
 113
 114   # Failed to find it. Download it!
 115   return _DownloadGsutil()
 116
 117
 118 def SupportsProdaccess(gsutil_path):
 119   with open(gsutil_path, 'r') as gsutil:
 120     return 'prodaccess' in gsutil.read()
 121
 122
 123 def _RunCommand(args):
 124   gsutil_path = FindGsutil()
 125
 126   # On cros device, as telemetry is running as root, home will be set to /root/,
 127   # which is not writable. gsutil will attempt to create a download tracker dir
 128   # in home dir and fail. To avoid this, override HOME dir to something writable
 129   # when running on cros device.
 130   #
 131   # TODO(tbarzic): Figure out a better way to handle gsutil on cros.
 132   #     http://crbug.com/386416, http://crbug.com/359293.
 133   gsutil_env = None
 134   if util.IsRunningOnCrosDevice():
 135     gsutil_env = os.environ.copy()
 136     gsutil_env['HOME'] = _CROS_GSUTIL_HOME_WAR
 137
 138   if os.name == 'nt':
 139     # If Windows, prepend python. Python scripts aren't directly executable.
 140     args = [sys.executable, gsutil_path] + args
 141   else:
 142     # Don't do it on POSIX, in case someone is using a shell script to redirect.
 143     args = [gsutil_path] + args
 144
 145   gsutil = subprocess.Popen(args, stdout=subprocess.PIPE,
 146                             stderr=subprocess.PIPE, env=gsutil_env)
 147   stdout, stderr = gsutil.communicate()
 148
 149   if gsutil.returncode:
 150     if stderr.startswith((
 151         'You are attempting to access protected data with no configured',
 152         'Failure: No handler was ready to authenticate.')):
 153       raise CredentialsError(gsutil_path)
 154     if ('status=403' in stderr or 'status 403' in stderr or
 155         '403 Forbidden' in stderr):
 156       raise PermissionError(gsutil_path)
 157     if (stderr.startswith('InvalidUriError') or 'No such object' in stderr or
 158         'No URLs matched' in stderr or 'One or more URLs matched no' in stderr):
 159       raise NotFoundError(stderr)
 160     if '500 Internal Server Error' in stderr:
 161       raise ServerError(stderr)
 162     raise CloudStorageError(stderr)
 163
 164   return stdout
 165
 166
 167 def List(bucket):
 168   query = 'gs://%s/' % bucket
 169   stdout = _RunCommand(['ls', query])
 170   return [url[len(query):] for url in stdout.splitlines()]
 171
 172
 173 def Exists(bucket, remote_path):
 174   try:
 175     _RunCommand(['ls', 'gs://%s/%s' % (bucket, remote_path)])
 176     return True
 177   except NotFoundError:
 178     return False
 179
 180
 181 def Move(bucket1, bucket2, remote_path):
 182   url1 = 'gs://%s/%s' % (bucket1, remote_path)
 183   url2 = 'gs://%s/%s' % (bucket2, remote_path)
 184   logging.info('Moving %s to %s' % (url1, url2))
 185   _RunCommand(['mv', url1, url2])
 186
 187
 188 def Copy(bucket_from, bucket_to, remote_path_from, remote_path_to):
 189   """Copy a file from one location in CloudStorage to another.
 190
 191   Args:
 192       bucket_from: The cloud storage bucket where the file is currently located.
 193       bucket_to: The cloud storage bucket it is being copied to.
 194       remote_path_from: The file path where the file is located in bucket_from.
 195       remote_path_to: The file path it is being copied to in bucket_to.
 196
 197   It should: cause no changes locally or to the starting file, and will
 198   overwrite any existing files in the destination location.
 199   """
 200   url1 = 'gs://%s/%s' % (bucket_from, remote_path_from)
 201   url2 = 'gs://%s/%s' % (bucket_to, remote_path_to)
 202   logging.info('Copying %s to %s' % (url1, url2))
 203   _RunCommand(['cp', url1, url2])
 204
 205
 206 def Delete(bucket, remote_path):
 207   url = 'gs://%s/%s' % (bucket, remote_path)
 208   logging.info('Deleting %s' % url)
 209   _RunCommand(['rm', url])
 210
 211
 212 def Get(bucket, remote_path, local_path):
 213   url = 'gs://%s/%s' % (bucket, remote_path)
 214   logging.info('Downloading %s to %s' % (url, local_path))
 215   try:
 216     _RunCommand(['cp', url, local_path])
 217   except ServerError:
 218     logging.info('Cloud Storage server error, retrying download')
 219     _RunCommand(['cp', url, local_path])
 220
 221
 222 def Insert(bucket, remote_path, local_path, publicly_readable=False):
 223   """ Upload file in |local_path| to cloud storage.
 224   Args:
 225     bucket: the google cloud storage bucket name.
 226     remote_path: the remote file path in |bucket|.
 227     local_path: path of the local file to be uploaded.
 228     publicly_readable: whether the uploaded file has publicly readable
 229     permission.
 230
 231   Returns:
 232     The url where the file is uploaded to.
 233   """
 234   url = 'gs://%s/%s' % (bucket, remote_path)
 235   command_and_args = ['cp']
 236   extra_info = ''
 237   if publicly_readable:
 238     command_and_args += ['-a', 'public-read']
 239     extra_info = ' (publicly readable)'
 240   command_and_args += [local_path, url]
 241   logging.info('Uploading %s to %s%s' % (local_path, url, extra_info))
 242   _RunCommand(command_and_args)
 243   return 'https://console.developers.google.com/m/cloudstorage/b/%s/o/%s' % (
 244       bucket, remote_path)
 245
 246
 247 def GetIfChanged(file_path, bucket):
 248   """Gets the file at file_path if it has a hash file that doesn't match or
 249   if there is no local copy of file_path, but there is a hash file for it.
 250
 251   Returns:
 252     True if the binary was changed.
 253   Raises:
 254     CredentialsError if the user has no configured credentials.
 255     PermissionError if the user does not have permission to access the bucket.
 256     NotFoundError if the file is not in the given bucket in cloud_storage.
 257   """
 258   hash_path = file_path + '.sha1'
 259   if not os.path.exists(hash_path):
 260     logging.warning('Hash file not found: %s' % hash_path)
 261     return False
 262
 263   expected_hash = ReadHash(hash_path)
 264   if os.path.exists(file_path) and CalculateHash(file_path) == expected_hash:
 265     return False
 266
 267   Get(bucket, expected_hash, file_path)
 268   return True
 269
 270 # TODO(aiolos): remove @decorators.Cache for http://crbug.com/459787
 271 @decorators.Cache
 272 def GetFilesInDirectoryIfChanged(directory, bucket):
 273   """ Scan the directory for .sha1 files, and download them from the given
 274   bucket in cloud storage if the local and remote hash don't match or
 275   there is no local copy.
 276   """
 277   if not os.path.isdir(directory):
 278     raise ValueError('Must provide a valid directory.')
 279   # Don't allow the root directory to be a serving_dir.
 280   if directory == os.path.abspath(os.sep):
 281     raise ValueError('Trying to serve root directory from HTTP server.')
 282   for dirpath, _, filenames in os.walk(directory):
 283     for filename in filenames:
 284       path_name, extension = os.path.splitext(
 285           os.path.join(dirpath, filename))
 286       if extension != '.sha1':
 287         continue
 288       GetIfChanged(path_name, bucket)
 289
 290 def CalculateHash(file_path):
 291   """Calculates and returns the hash of the file at file_path."""
 292   sha1 = hashlib.sha1()
 293   with open(file_path, 'rb') as f:
 294     while True:
 295       # Read in 1mb chunks, so it doesn't all have to be loaded into memory.
 296       chunk = f.read(1024*1024)
 297       if not chunk:
 298         break
 299       sha1.update(chunk)
 300   return sha1.hexdigest()
 301
 302
 303 def ReadHash(hash_path):
 304   with open(hash_path, 'rb') as f:
 305     return f.read(1024).rstrip()