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