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