1 #!/usr/bin/env python 2 # Copyright 2014 The Chromium Authors. All rights reserved. 3 # Use of this source code is governed by a BSD-style license that can be 4 # found in the LICENSE file. 5 6 import argparse 7 import logging 8 import os 9 import subprocess 10 import sys 11 12 from telemetry.core import command_line 13 from telemetry.page import cloud_storage 14 15 16 BUCKET_ALIASES = { 17 'public': cloud_storage.PUBLIC_BUCKET, 18 'partner': cloud_storage.PARTNER_BUCKET, 19 'google-only': cloud_storage.INTERNAL_BUCKET, 20 } 21 BUCKETS = {bucket: easy_bucket_name for easy_bucket_name, bucket 22 in BUCKET_ALIASES.iteritems()} 23 24 25 def _GetPaths(path): 26 root, ext = os.path.splitext(path) 27 if ext == '.sha1': 28 file_path = root 29 hash_path = path 30 else: 31 file_path = path 32 hash_path = path + '.sha1' 33 return file_path, hash_path 34 35 36 def _FindFilesInCloudStorage(files): 37 """Returns a dict of all files and which buckets they're in.""" 38 # Preprocessing: get the contents of all buckets. 39 bucket_contents = {} 40 for bucket in BUCKETS: 41 try: 42 bucket_contents[bucket] = cloud_storage.List(bucket) 43 except (cloud_storage.PermissionError, cloud_storage.CredentialsError): 44 pass 45 46 # Check if each file is in the bucket contents. 47 file_buckets = {} 48 for path in files: 49 file_path, hash_path = _GetPaths(path) 50 51 if file_path in file_buckets: 52 # Ignore duplicates, if both data and sha1 file were in the file list. 53 continue 54 if not os.path.exists(hash_path): 55 # Probably got some non-Cloud Storage files in the file list. Ignore. 56 continue 57 58 file_hash = cloud_storage.ReadHash(hash_path) 59 file_buckets[file_path] = [] 60 for bucket in BUCKETS: 61 if bucket in bucket_contents and file_hash in bucket_contents[bucket]: 62 file_buckets[file_path].append(bucket) 63 64 return file_buckets 65 66 67 class Ls(command_line.Command): 68 """List which bucket each file is in.""" 69 70 @classmethod 71 def AddCommandLineArgs(cls, parser): 72 parser.add_argument('-r', '--recursive', action='store_true') 73 parser.add_argument('paths', nargs='+') 74 75 @classmethod 76 def ProcessCommandLineArgs(cls, parser, args): 77 for path in args.paths: 78 if not os.path.exists(path): 79 parser.error('Path not found: %s' % path) 80 81 def Run(self, args): 82 def GetFilesInPaths(paths, recursive): 83 """If path is a dir, yields all files in path, otherwise just yields path. 84 85 If recursive is true, walks subdirectories recursively.""" 86 for path in paths: 87 if not os.path.isdir(path): 88 yield path 89 continue 90 91 if recursive: 92 for root, _, filenames in os.walk(path): 93 for filename in filenames: 94 yield os.path.join(root, filename) 95 else: 96 for filename in os.listdir(path): 97 yield os.path.join(path, filename) 98 99 files = _FindFilesInCloudStorage(GetFilesInPaths(args.paths, args.recursive)) 100 101 if not files: 102 print 'No files in Cloud Storage.' 103 return 104 105 for file_path, buckets in sorted(files.iteritems()): 106 if buckets: 107 buckets = [BUCKETS[bucket] for bucket in buckets] 108 print '%-11s %s' % (','.join(buckets), file_path) 109 else: 110 print '%-11s %s' % ('not found', file_path) 111 112 113 class Mv(command_line.Command): 114 """Move files to the given bucket.""" 115 116 @classmethod 117 def AddCommandLineArgs(cls, parser): 118 parser.add_argument('files', nargs='+') 119 parser.add_argument('bucket', choices=BUCKET_ALIASES) 120 121 @classmethod 122 def ProcessCommandLineArgs(cls, parser, args): 123 args.bucket = BUCKET_ALIASES[args.bucket] 124 125 def Run(self, args): 126 files = _FindFilesInCloudStorage(args.files) 127 128 for file_path, buckets in sorted(files.iteritems()): 129 if not buckets: 130 raise IOError('%s not found in Cloud Storage.' % file_path) 131 132 for file_path, buckets in sorted(files.iteritems()): 133 if args.bucket in buckets: 134 buckets.remove(args.bucket) 135 if not buckets: 136 logging.info('Skipping %s, no action needed.' % file_path) 137 continue 138 139 # Move to the target bucket. 140 file_hash = cloud_storage.ReadHash(file_path + '.sha1') 141 cloud_storage.Move(buckets.pop(), args.bucket, file_hash) 142 143 # Delete all additional copies. 144 for bucket in buckets: 145 cloud_storage.Delete(bucket, file_hash) 146 147 148 class Rm(command_line.Command): 149 """Remove files from Cloud Storage.""" 150 151 @classmethod 152 def AddCommandLineArgs(cls, parser): 153 parser.add_argument('files', nargs='+') 154 155 def Run(self, args): 156 files = _FindFilesInCloudStorage(args.files) 157 for file_path, buckets in sorted(files.iteritems()): 158 file_hash = cloud_storage.ReadHash(file_path + '.sha1') 159 for bucket in buckets: 160 cloud_storage.Delete(bucket, file_hash) 161 162 163 class Upload(command_line.Command): 164 """Upload files to Cloud Storage.""" 165 166 @classmethod 167 def AddCommandLineArgs(cls, parser): 168 parser.add_argument('files', nargs='+') 169 parser.add_argument('bucket', choices=BUCKET_ALIASES) 170 171 @classmethod 172 def ProcessCommandLineArgs(cls, parser, args): 173 args.bucket = BUCKET_ALIASES[args.bucket] 174 175 for path in args.files: 176 if not os.path.exists(path): 177 parser.error('File not found: %s' % path) 178 179 def Run(self, args): 180 for file_path in args.files: 181 file_hash = cloud_storage.CalculateHash(file_path) 182 183 # Create or update the hash file. 184 hash_path = file_path + '.sha1' 185 with open(hash_path, 'wb') as f: 186 f.write(file_hash) 187 f.flush() 188 189 # Add the data to Cloud Storage. 190 cloud_storage.Insert(args.bucket, file_hash, file_path) 191 192 # Add the hash file to the branch, for convenience. :) 193 subprocess.call(['git', 'add', hash_path]) 194 195 196 class CloudStorageCommand(command_line.SubcommandCommand): 197 commands = (Ls, Mv, Rm, Upload) 198 199 200 if __name__ == '__main__': 201 logging.getLogger().setLevel(logging.INFO) 202 sys.exit(CloudStorageCommand.main()) 203
