1# Copyright 2013 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 5import json 6import logging 7import os 8import shutil 9import tempfile 10import time 11 12from py_utils import cloud_storage # pylint: disable=import-error 13 14 15_DEFAULT_PLATFORM = 'DEFAULT' 16_ALL_PLATFORMS = ['mac', 'linux', 'android', 'win', _DEFAULT_PLATFORM] 17 18 19def AssertValidCloudStorageBucket(bucket): 20 is_valid = bucket in (None, 21 cloud_storage.PUBLIC_BUCKET, 22 cloud_storage.PARTNER_BUCKET, 23 cloud_storage.INTERNAL_BUCKET) 24 if not is_valid: 25 raise ValueError("Cloud storage privacy bucket %s is invalid" % bucket) 26 27 28class WprArchiveInfo(object): 29 def __init__(self, file_path, data, bucket): 30 AssertValidCloudStorageBucket(bucket) 31 self._file_path = file_path 32 self._base_dir = os.path.dirname(file_path) 33 self._data = data 34 self._bucket = bucket 35 self.temp_target_wpr_file_path = None 36 # Ensure directory exists. 37 if not os.path.exists(self._base_dir): 38 os.makedirs(self._base_dir) 39 40 assert data.get('platform_specific', False), ( 41 'Detected old version of archive info json file. Please update to new ' 42 'version.') 43 44 self._story_name_to_wpr_file = data['archives'] 45 46 @classmethod 47 def FromFile(cls, file_path, bucket): 48 """ Generates an archive_info instance with the given json file. """ 49 if os.path.exists(file_path): 50 with open(file_path, 'r') as f: 51 data = json.load(f) 52 return cls(file_path, data, bucket) 53 return cls(file_path, {'archives': {}, 'platform_specific': True}, bucket) 54 55 def DownloadArchivesIfNeeded(self, target_platforms=None, story_names=None): 56 """Downloads archives iff the Archive has a bucket parameter and the user 57 has permission to access the bucket. 58 59 Raises cloud storage Permissions or Credentials error when there is no 60 local copy of the archive and the user doesn't have permission to access 61 the archive's bucket. 62 63 Warns when a bucket is not specified or when the user doesn't have 64 permission to access the archive's bucket but a local copy of the archive 65 exists. 66 67 Args: 68 target_platform: only downloads archives for these platforms 69 story_names: only downloads archives for these story names 70 """ 71 logging.info('Downloading WPR archives. This can take a long time.') 72 start_time = time.time() 73 # If no target platform is set, download all platforms. 74 if target_platforms is None: 75 target_platforms = _ALL_PLATFORMS 76 else: 77 assert isinstance(target_platforms, list), 'Must pass platforms as a list' 78 target_platforms = target_platforms + [_DEFAULT_PLATFORM] 79 # Download all .wprgo files. 80 if not self._bucket: 81 logging.warning('Story set in %s has no bucket specified, and ' 82 'cannot be downloaded from cloud_storage.', ) 83 return 84 assert 'archives' in self._data, ("Invalid data format in %s. 'archives' " 85 "field is needed" % self._file_path) 86 87 def download_if_needed(path): 88 try: 89 cloud_storage.GetIfChanged(path, self._bucket) 90 except (cloud_storage.CredentialsError, cloud_storage.PermissionError): 91 if os.path.exists(path): 92 # If the archive exists, assume the user recorded their own and warn 93 # them that they do not have the proper credentials to download. 94 logging.warning('Need credentials to update WPR archive: %s', path) 95 else: 96 logging.error("You either aren't authenticated or don't have " 97 "permission to use the archives for this page set." 98 "\nYou may need to run gsutil config." 99 "\nYou can find instructions for gsutil config at: " 100 "http://www.chromium.org/developers/telemetry/" 101 "upload_to_cloud_storage") 102 raise 103 104 try: 105 story_archives = self._data['archives'] 106 download_names = set(story_archives.iterkeys()) 107 if story_names is not None: 108 download_names.intersection_update(story_names) 109 for story_name in download_names: 110 for target_platform in target_platforms: 111 if story_archives[story_name].get(target_platform): 112 archive_path = self._WprFileNameToPath( 113 story_archives[story_name][target_platform]) 114 download_if_needed(archive_path) 115 finally: 116 logging.info('All WPR archives are downloaded, took %s seconds.', 117 time.time() - start_time) 118 119 def WprFilePathForStory(self, story, target_platform=_DEFAULT_PLATFORM): 120 if self.temp_target_wpr_file_path: 121 return self.temp_target_wpr_file_path 122 123 wpr_file = self._story_name_to_wpr_file.get(story.name, None) 124 if wpr_file: 125 if target_platform in wpr_file: 126 return self._WprFileNameToPath(wpr_file[target_platform]) 127 return self._WprFileNameToPath(wpr_file[_DEFAULT_PLATFORM]) 128 return None 129 130 def AddNewTemporaryRecording(self, temp_wpr_file_path=None): 131 if temp_wpr_file_path is None: 132 temp_wpr_file_handle, temp_wpr_file_path = tempfile.mkstemp() 133 os.close(temp_wpr_file_handle) 134 self.temp_target_wpr_file_path = temp_wpr_file_path 135 136 def AddRecordedStories(self, stories, upload_to_cloud_storage=False, 137 target_platform=_DEFAULT_PLATFORM): 138 if not stories: 139 os.remove(self.temp_target_wpr_file_path) 140 return 141 142 target_wpr_file_hash = cloud_storage.CalculateHash( 143 self.temp_target_wpr_file_path) 144 (target_wpr_file, target_wpr_file_path) = self._NextWprFileName( 145 target_wpr_file_hash) 146 for story in stories: 147 # Check to see if the platform has been manually overrided. 148 if not story.platform_specific: 149 current_target_platform = _DEFAULT_PLATFORM 150 else: 151 current_target_platform = target_platform 152 self._SetWprFileForStory( 153 story.name, target_wpr_file, current_target_platform) 154 shutil.move(self.temp_target_wpr_file_path, target_wpr_file_path) 155 156 # Update the hash file. 157 with open(target_wpr_file_path + '.sha1', 'wb') as f: 158 f.write(target_wpr_file_hash) 159 f.flush() 160 161 self._WriteToFile() 162 163 # Upload to cloud storage 164 if upload_to_cloud_storage: 165 if not self._bucket: 166 logging.warning('StorySet must have bucket specified to upload ' 167 'stories to cloud storage.') 168 return 169 try: 170 cloud_storage.Insert(self._bucket, target_wpr_file_hash, 171 target_wpr_file_path) 172 except cloud_storage.CloudStorageError, e: 173 logging.warning('Failed to upload wpr file %s to cloud storage. ' 174 'Error:%s' % target_wpr_file_path, e) 175 176 def _WriteToFile(self): 177 """Writes the metadata into the file passed as constructor parameter.""" 178 metadata = dict() 179 metadata['description'] = ( 180 'Describes the Web Page Replay archives for a story set. ' 181 'Don\'t edit by hand! Use record_wpr for updating.') 182 metadata['archives'] = self._story_name_to_wpr_file.copy() 183 metadata['platform_specific'] = True 184 185 with open(self._file_path, 'w') as f: 186 json.dump(metadata, f, indent=4, sort_keys=True, separators=(',', ': ')) 187 f.flush() 188 189 def _WprFileNameToPath(self, wpr_file): 190 return os.path.abspath(os.path.join(self._base_dir, wpr_file)) 191 192 def _NextWprFileName(self, file_hash): 193 """Creates a new file name for a wpr archive file.""" 194 base = os.path.splitext(os.path.basename(self._file_path))[0] 195 new_filename = '%s_%s.%s' % (base, file_hash[:10], 'wprgo') 196 return new_filename, self._WprFileNameToPath(new_filename) 197 198 def _SetWprFileForStory(self, story_name, wpr_file, target_platform): 199 """For modifying the metadata when we're going to record a new archive.""" 200 if story_name not in self._data['archives']: 201 # If there is no other recording we want the first to be the default 202 # until a new default is recorded. 203 self._data['archives'][story_name] = {_DEFAULT_PLATFORM: wpr_file} 204 self._data['archives'][story_name][target_platform] = wpr_file 205