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