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