tools/perf/profile_creators/extension_profile_extender.py

   1 # Copyright 2015 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 import atexit
   6 import json
   7 import logging
   8 import os
   9 import shutil
  10 import sys
  11 import time
  12 import zipfile
  13
  14 if sys.platform == 'win32':
  15   import _winreg as winreg # pylint: disable=import-error
  16
  17 from catapult_base import cloud_storage
  18 from profile_creators import profile_extender
  19 from telemetry.core import exceptions
  20
  21
  22 # Remote target upload directory in cloud storage for extensions.
  23 REMOTE_DIR = 'extension_set'
  24
  25 # Target zip file.
  26 ZIP_NAME = 'extensions.zip'
  27
  28
  29 class InvalidExtensionArchiveError(exceptions.Error):
  30   """Exception thrown when remote archive is invalid or malformed.
  31
  32   Remote archive should be located at REMOTE_DIR/ZIP_NAME. Upon failure,
  33   prompts user to update remote archive using update_remote_extensions
  34   script.
  35   """
  36
  37   def __init__(self, msg=''):
  38     msg += ('\nTry running\n'
  39             '\tpython update_remote_extensions.py -e extension_set.csv\n'
  40             'in src/tools/perf/profile_creator subdirectory.')
  41     super(InvalidExtensionArchiveError, self).__init__(msg)
  42
  43
  44 class ExtensionProfileExtender(profile_extender.ProfileExtender):
  45   """Creates a profile with many extensions."""
  46
  47   def __init__(self, finder_options):
  48     super(ExtensionProfileExtender, self).__init__(finder_options)
  49     self._extensions = []
  50     finder_options.browser_options.disable_default_apps = False
  51     finder_options.browser_options.AppendExtraBrowserArgs(
  52         '--prompt-for-external-extensions=0')
  53
  54   def Run(self):
  55     """Superclass override."""
  56     # Download extensions from cloud and force-install extensions into profile.
  57     local_extensions_dir = os.path.join(self.profile_path,
  58                                         'external_extensions_crx')
  59     self._DownloadRemoteExtensions(cloud_storage.PARTNER_BUCKET,
  60                                    local_extensions_dir)
  61     atexit.register(self._CleanUpExtensions)
  62     self._LoadExtensions(local_extensions_dir, self.profile_path)
  63     try:
  64       self.SetUpBrowser()
  65       self._WaitForExtensionsToLoad()
  66     finally:
  67       self.TearDownBrowser()
  68
  69   def _DownloadRemoteExtensions(self, remote_bucket, local_extensions_dir):
  70     """Downloads and unzips archive of common extensions to disk.
  71
  72     Args:
  73         remote_bucket: bucket to download remote archive from.
  74         local_extensions_dir: destination extensions directory.
  75
  76     Raises:
  77         InvalidExtensionArchiveError if remote archive is not found.
  78     """
  79     # Force Unix directory separator for remote path.
  80     remote_zip_path = '%s/%s' % (REMOTE_DIR, ZIP_NAME)
  81     local_zip_path = os.path.join(local_extensions_dir, ZIP_NAME)
  82     try:
  83       cloud_storage.Get(remote_bucket, remote_zip_path, local_zip_path)
  84     except:
  85       raise InvalidExtensionArchiveError('Can\'t find archive at gs://%s/%s..'
  86                                          % (remote_bucket, remote_zip_path))
  87     try:
  88       with zipfile.ZipFile(local_zip_path, 'r') as extensions_zip:
  89         extensions_zip.extractall(local_extensions_dir)
  90     finally:
  91       os.remove(local_zip_path)
  92
  93   def _GetExtensionInfoFromCrx(self, crx_file):
  94     """Retrieves version + name of extension from CRX archive."""
  95     with zipfile.ZipFile(crx_file, 'r') as crx_zip:
  96       manifest_contents = crx_zip.read('manifest.json')
  97       decoded_manifest = json.loads(manifest_contents)
  98       crx_version = decoded_manifest['version']
  99       extension_name = decoded_manifest['name']
 100     return (crx_version, extension_name)
 101
 102   def _LoadExtensions(self, local_extensions_dir, profile_dir):
 103     """Loads extensions in _local_extensions_dir into user profile.
 104
 105     Extensions are loaded according to platform specifications at
 106     https://developer.chrome.com/extensions/external_extensions.html
 107
 108     Args:
 109         local_extensions_dir: directory containing CRX files.
 110         profile_dir: target profile directory for the extensions.
 111
 112     Raises:
 113         InvalidExtensionArchiveError if archive contains a non-CRX file.
 114     """
 115     ext_files = os.listdir(local_extensions_dir)
 116     external_ext_dir = os.path.join(profile_dir, 'External Extensions')
 117     os.makedirs(external_ext_dir)
 118     for ext_file in ext_files:
 119       ext_path = os.path.join(local_extensions_dir, ext_file)
 120       if not ext_file.endswith('.crx'):
 121         raise InvalidExtensionArchiveError('Archive contains non-crx file %s.'
 122                                            % ext_file)
 123       (version, name) = self._GetExtensionInfoFromCrx(ext_path)
 124       ext_id = os.path.splitext(os.path.basename(ext_path))[0]
 125       extension_info = {
 126         'extension_id': ext_id,
 127         'external_crx': ext_path,
 128         'external_version': version,
 129         '_comment': name
 130       }
 131       # Platform-specific external extension installation
 132       if self.os_name == 'win': # Windows
 133         key_path = 'Software\\Google\\Chrome\\Extensions\\%s' % ext_id
 134         self._WriteRegistryValue(key_path, 'Path', ext_path)
 135         self._WriteRegistryValue(key_path, 'Version', version)
 136       else:
 137         extension_json_path = os.path.join(external_ext_dir, '%s.json' % ext_id)
 138         with open(extension_json_path, 'w') as f:
 139           f.write(json.dumps(extension_info))
 140       self._extensions.append(ext_id)
 141
 142   def _WriteRegistryValue(self, key_path, name, value):
 143     """Writes (or overwrites) registry value specified to HKCU\\key_path."""
 144     with winreg.CreateKey(winreg.HKEY_CURRENT_USER, key_path) as key:
 145       try: # Does registry value already exist?
 146         path_value = winreg.QueryValueEx(key, name)
 147         if path_value != value:
 148           logging.warning(
 149               'Overwriting registry value %s\\%s:'
 150               '\n%s with %s' % (key_path, name, path_value, value))
 151       except OSError:
 152         pass
 153       winreg.SetValueEx(key, name, 0, winreg.REG_SZ, value)
 154
 155   def _CleanUpExtensions(self):
 156     """Cleans up registry keys or JSON files used to install extensions."""
 157     if self.os_name == 'win':
 158       for ext_id in self._extensions:
 159         winreg.DeleteKey(winreg.HKEY_CURRENT_USER,
 160             'Software\\Google\\Chrome\\Extensions\\%s' % ext_id)
 161     else:
 162       to_remove = os.path.join(self.profile_path, 'External Extensions')
 163       if os.path.exists(to_remove):
 164         shutil.rmtree(to_remove)
 165
 166   def _WaitForExtensionsToLoad(self):
 167     """Stall until browser has finished installing/loading all extensions."""
 168     unloaded_extensions = set(self._extensions)
 169     while unloaded_extensions:
 170       loaded_extensions = set([key.extension_id for key in
 171                               self.browser.extensions.keys()])
 172       unloaded_extensions = unloaded_extensions - loaded_extensions
 173       # There's no event signalling when browser finishes installing
 174       # or loading an extension so re-check every 5 seconds.
 175       time.sleep(5)