text: # Copyright 2013 The Chromium Authors. All rights reserved.
  # Use of this source code is governed by a BSD-style license that can be
  # found in the LICENSE file.
  
  import json
  import logging
  import os
  import re
  import shutil
  import tempfile
  
  from py_utils import cloud_storage  # pylint: disable=import-error
  
  
  _DEFAULT_PLATFORM = 'DEFAULT'
  _ALL_PLATFORMS = ['mac', 'linux', 'android', 'win', _DEFAULT_PLATFORM]
  
  
  # TODO(rnephew): Move this file to archive_info.py when it is the default.
  # These changes were chosen to not be merged with archive_info.py directly
  # because all code in archive_info.py will be deleted once everything is
  # switched to using this version.
  def AssertValidCloudStorageBucket(bucket):
    is_valid = bucket in (None,
                          cloud_storage.PUBLIC_BUCKET,
                          cloud_storage.PARTNER_BUCKET,
                          cloud_storage.INTERNAL_BUCKET)
    if not is_valid:
      raise ValueError("Cloud storage privacy bucket %s is invalid" % bucket)
  
  
  class WprArchiveInfo(object):
    def __init__(self, file_path, data, bucket):
      AssertValidCloudStorageBucket(bucket)
      self._file_path = file_path
      self._base_dir = os.path.dirname(file_path)
      self._data = data
      self._bucket = bucket
      self.temp_target_wpr_file_path = None
      # Ensure directory exists.
      if not os.path.exists(self._base_dir):
        os.makedirs(self._base_dir)
  
      assert data.get('platform_specific', False), (
          'archive_info2.WprArchiveInfo is exclusively for platform specific WPR '
          'archives. Platform agnostic archives should use '
          'archive_info.WprArchiveInfo archive_info2 will become archive_info '
          'once all data files are switched.')
  
      self._story_name_to_wpr_file = data['archives']
  
    @classmethod
    def FromFile(cls, file_path, bucket):
      """ Generates an archive_info2 instance with the given json file. """
      if os.path.exists(file_path):
        with open(file_path, 'r') as f:
          data = json.load(f)
          return cls(file_path, data, bucket)
      return cls(file_path, {'archives': {}}, bucket)
  
    def DownloadArchivesIfNeeded(self, target_platforms=None):
      """Downloads archives iff the Archive has a bucket parameter and the user
      has permission to access the bucket.
  
      Raises cloud storage Permissions or Credentials error when there is no
      local copy of the archive and the user doesn't have permission to access
      the archive's bucket.
  
      Warns when a bucket is not specified or when the user doesn't have
      permission to access the archive's bucket but a local copy of the archive
      exists.
      """
      # If no target platform is set, download all platforms.
      if target_platforms is None:
        target_platforms = _ALL_PLATFORMS
      else:
        assert isinstance(target_platforms, list), 'Must pass platforms as a list'
        target_platforms = target_platforms + [_DEFAULT_PLATFORM]
      # Download all .wpr files.
      if not self._bucket:
        logging.warning('Story set in %s has no bucket specified, and '
                        'cannot be downloaded from cloud_storage.', )
        return
      assert 'archives' in self._data, ("Invalid data format in %s. 'archives' "
                                        "field is needed" % self._file_path)
  
      def download_if_needed(path):
        try:
          cloud_storage.GetIfChanged(path, self._bucket)
        except (cloud_storage.CredentialsError, cloud_storage.PermissionError):
          if os.path.exists(path):
            # If the archive exists, assume the user recorded their own and warn
            # them that they do not have the proper credentials to download.
            logging.warning('Need credentials to update WPR archive: %s', path)
          else:
            logging.error("You either aren't authenticated or don't have "
                          "permission to use the archives for this page set."
                          "\nYou may need to run gsutil config."
                          "\nYou can find instructions for gsutil config at: "
                          "http://www.chromium.org/developers/telemetry/"
                          "upload_to_cloud_storage")
            raise
  
      story_archives = self._data['archives']
      for story in story_archives:
        for target_platform in target_platforms:
          if story_archives[story].get(target_platform):
            archive_path = self._WprFileNameToPath(
                story_archives[story][target_platform])
            download_if_needed(archive_path)
  
    def WprFilePathForStory(self, story, target_platform=_DEFAULT_PLATFORM):
      if self.temp_target_wpr_file_path:
        return self.temp_target_wpr_file_path
  
      wpr_file = self._story_name_to_wpr_file.get(story.display_name, None)
      if wpr_file is None and hasattr(story, 'url'):
        # Some old pages always use the URL to identify a page rather than the
        # display_name, so try to look for that.
        wpr_file = self._story_name_to_wpr_file.get(story.url, None)
      if wpr_file:
        if target_platform in wpr_file:
          return self._WprFileNameToPath(wpr_file[target_platform])
        return self._WprFileNameToPath(wpr_file[_DEFAULT_PLATFORM])
      return None
  
    def AddNewTemporaryRecording(self, temp_wpr_file_path=None):
      if temp_wpr_file_path is None:
        temp_wpr_file_handle, temp_wpr_file_path = tempfile.mkstemp()
        os.close(temp_wpr_file_handle)
      self.temp_target_wpr_file_path = temp_wpr_file_path
  
    def AddRecordedStories(self, stories, upload_to_cloud_storage=False,
                           target_platform=_DEFAULT_PLATFORM):
      if not stories:
        os.remove(self.temp_target_wpr_file_path)
        return
  
      (target_wpr_file, target_wpr_file_path) = self._NextWprFileName()
      for story in stories:
        # Check to see if the platform has been manually overrided.
        if not story.platform_specific:
          current_target_platform = _DEFAULT_PLATFORM
        else:
          current_target_platform = target_platform
        self._SetWprFileForStory(
            story.display_name, target_wpr_file, current_target_platform)
      shutil.move(self.temp_target_wpr_file_path, target_wpr_file_path)
  
      # Update the hash file.
      target_wpr_file_hash = cloud_storage.CalculateHash(target_wpr_file_path)
      with open(target_wpr_file_path + '.sha1', 'wb') as f:
        f.write(target_wpr_file_hash)
        f.flush()
  
      self._WriteToFile()
  
      # Upload to cloud storage
      if upload_to_cloud_storage:
        if not self._bucket:
          logging.warning('StorySet must have bucket specified to upload '
                          'stories to cloud storage.')
          return
        try:
          cloud_storage.Insert(self._bucket, target_wpr_file_hash,
                               target_wpr_file_path)
        except cloud_storage.CloudStorageError, e:
          logging.warning('Failed to upload wpr file %s to cloud storage. '
                          'Error:%s' % target_wpr_file_path, e)
  
    def _WriteToFile(self):
      """Writes the metadata into the file passed as constructor parameter."""
      metadata = dict()
      metadata['description'] = (
          'Describes the Web Page Replay archives for a story set. '
          'Don\'t edit by hand! Use record_wpr for updating.')
      metadata['archives'] = self._story_name_to_wpr_file.copy()
      metadata['platform_specific'] = True
  
      with open(self._file_path, 'w') as f:
        json.dump(metadata, f, indent=4, sort_keys=True, separators=(',', ': '))
        f.flush()
  
    def _WprFileNameToPath(self, wpr_file):
      return os.path.abspath(os.path.join(self._base_dir, wpr_file))
  
    def _NextWprFileName(self):
      """Creates a new file name for a wpr archive file."""
      # The names are of the format "some_thing_number.wpr". Read the numbers.
      highest_number = -1
      base = None
      wpr_files = []
      for story in self._data['archives']:
        for p in self._data['archives'][story]:
          wpr_files.append(self._data['archives'][story][p])
  
      for wpr_file in wpr_files:
        match = re.match(r'(?P<BASE>.*)_(?P<NUMBER>[0-9]+)\.wpr', wpr_file)
        if not match:
          raise Exception('Illegal wpr file name ' + wpr_file)
        highest_number = max(int(match.groupdict()['NUMBER']), highest_number)
        if base and match.groupdict()['BASE'] != base:
          raise Exception('Illegal wpr file name ' + wpr_file +
                          ', doesn\'t begin with ' + base)
        base = match.groupdict()['BASE']
      if not base:
        # If we're creating a completely new info file, use the base name of the
        # story set file.
        base = os.path.splitext(os.path.basename(self._file_path))[0]
      new_filename = '%s_%03d.wpr' % (base, highest_number + 1)
      return new_filename, self._WprFileNameToPath(new_filename)
  
    def _SetWprFileForStory(self, story_name, wpr_file, target_platform):
      """For modifying the metadata when we're going to record a new archive."""
      if story_name not in self._data['archives']:
        # If there is no other recording we want the first to be the default
        # until a new default is recorded.
        self._data['archives'][story_name] = {_DEFAULT_PLATFORM: wpr_file}
      self._data['archives'][story_name][target_platform] = wpr_file
  