tools/perf/profile_creators/update_remote_extensions.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 base64
   6 import csv
   7 import json
   8 import optparse
   9 import os
  10 import shutil
  11 import sys
  12 import tempfile
  13 import urllib2
  14 import zipfile
  15
  16 sys.path.append(os.path.join(os.path.dirname(__file__), os.pardir, os.pardir,
  17     'telemetry'))
  18
  19 from catapult_base import cloud_storage
  20 from telemetry.core import exceptions
  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 def _DownloadCrxFromCws(ext_id, dst):
  30   """Downloads CRX specified from Chrome Web Store.
  31
  32   Retrieves CRX (Chrome extension file) specified by ext_id from Chrome Web
  33   Store, into directory specified by dst.
  34
  35   Args:
  36       ext_id: id of extension to retrieve.
  37       dst: directory to download CRX into
  38
  39   Returns:
  40       Returns local path to downloaded CRX.
  41       If download fails, return None.
  42   """
  43   dst_path = os.path.join(dst, '%s.crx' % ext_id)
  44   cws_url = ('https://clients2.google.com/service/update2/crx?response='
  45              'redirect&prodversion=38.0&x=id%%3D%s%%26installsource%%3D'
  46              'ondemand%%26uc' % ext_id)
  47   response = urllib2.urlopen(cws_url)
  48   if response.getcode() is not 200:
  49     return None
  50   with open(dst_path, 'w') as f:
  51     f.write(response.read())
  52   return dst_path
  53
  54
  55 def _UpdateExtensionsInCloud(local_extensions_dir, extensions_csv, remote_dir):
  56   """Updates set of extensions in Cloud Storage from a CSV of extension ids.
  57
  58   From well-formatted CSV file containing some set of extensions
  59   (extensions_csv), download them, compress into archive, and update
  60   the remote extension archive under REMOTE_DIR in CHROME-PARTNER-TELEMETRY
  61   bucket. This script expects 2nd column of CSV file to contain extension ids.
  62
  63   Args:
  64       local_extensions_dir: directory to download CRX files into.
  65       extension_csv: CSV to pull extension_ids from.
  66       remote_dir: remote directory to put extension archive in cloud storage.
  67
  68   Raises:
  69       Exception if a CRX download fails.
  70   """
  71
  72   # Download CRX to temp files and compress into archive
  73   zip_path = os.path.join(local_extensions_dir, ZIP_NAME)
  74   extension_zip = zipfile.ZipFile(zip_path, 'w')
  75   update_csv = False
  76   extensions_info = []
  77   with open(extensions_csv, 'rb') as csv_file:
  78     reader = csv.reader(csv_file)
  79     # Stores comments (in case CSV needs to be updated/rewritten)
  80     # and skips header line.
  81     comments = []
  82     line = ','.join(reader.next())
  83     while line.startswith('#'):
  84       comments.append(line)
  85       line = ','.join(reader.next())
  86     # Extract info from CSV.
  87     for row in reader:
  88       extension_info = {
  89         'extension_name': row[0],
  90         'id': row[1],
  91         'hash': row[2],
  92         'version': row[3]
  93       }
  94
  95       print 'Fetching extension %s...' % extension_info['id']
  96       crx_path = _DownloadCrxFromCws(extension_info['id'], local_extensions_dir)
  97       if crx_path is None:
  98         raise exceptions.Error('\tCould not fetch %s.\n\n'
  99                                'If this extension dl consistently fails, '
 100                                'remove this entry from %s.'
 101                                % (extension_info['id'], extensions_csv))
 102       (new_hash, new_version) = _CrxHashIfChanged(crx_path, extension_info)
 103       if new_hash is not None:
 104         update_csv = True
 105         extension_info['hash'] =  new_hash
 106         extension_info['version'] = new_version
 107       extensions_info.append(extension_info)
 108       extension_zip.write(crx_path, arcname='%s.crx' % extension_info['id'])
 109   extension_zip.close()
 110
 111   if update_csv:
 112     print 'Updating CSV...'
 113     _UpdateCsv(comments, extensions_csv, extensions_info)
 114
 115   print 'Uploading extensions to cloud...'
 116   remote_zip_path = os.path.join(remote_dir, ZIP_NAME)
 117   cloud_storage.Insert(cloud_storage.PARTNER_BUCKET, remote_zip_path, zip_path)
 118
 119
 120 def _CrxHashIfChanged(crx_path, extension_info):
 121   """Checks whether downloaded Crx has been altered.
 122
 123   Compares stored hash with hash of downloaded Crx. If different, alerts user
 124   that CRX version has changed and will be updated in CSV file.
 125
 126   Args:
 127     crx_path: Path to downloaded CRX.
 128     extension_info: Info from CSV (including id and previous hash) about CRX.
 129
 130   Returns:
 131     New hash and version if extension differed. Otherwise, returns (None, None)
 132   """
 133   downloaded_hash = _Base64Hash(crx_path)
 134   new_version = _GetVersionFromCrx(crx_path)
 135   if downloaded_hash != extension_info['hash']:
 136     if new_version != extension_info['version']:
 137       ans = raw_input('\tWarning: Extension %s version from Web Store differs '
 138                       'from CSV version.\n\tIf continued, script will write '
 139                       'new hash and version to CSV.\n\tContinue? (y/n) '
 140                       % extension_info['id']).lower()
 141     else:
 142       raise exceptions.Error('Extension %s hash from Web Store differs from '
 143                              '\nhash stored in CSV, but versions are the same.')
 144     if not ans.startswith('y'):
 145       sys.exit('Web Store extension %s hash differs from hash in CSV.'
 146                % extension_info['id'])
 147     return (downloaded_hash, new_version)
 148   return (None, None)
 149
 150 def _UpdateCsv(comments, extensions_csv, extensions_info):
 151   """Updates CSV with information in extensions_info.
 152
 153   Original CSV is overwritten with updated information about each extension.
 154   Header comments from original CSV are preserved.
 155
 156   Args:
 157     comments: List containing lines of comments found in header of original CSV.
 158     extensions_csv: Path to CSV file.
 159     extensions_info: List of extension info to write to CSV. Each entry is
 160         a dict containing fields extension_name, id, hash, and version.
 161   """
 162   # Maintain pre-existing comments.
 163   with open(extensions_csv, 'w') as csv_file:
 164     csv_file.write('\n'.join(comments))
 165     csv_file.write('\n')
 166   with open(extensions_csv, 'a') as csv_file:
 167     writer = csv.DictWriter(
 168         csv_file, fieldnames=['extension_name', 'id', 'hash', 'version'])
 169     writer.writeheader()
 170     writer.writerows(extensions_info)
 171
 172 def _GetCsvFromArgs():
 173   """Parse options to retrieve name of CSV file."""
 174   parser = optparse.OptionParser()
 175   parser.add_option('-e', '--extension-csv', dest='extension_csv',
 176                    help='CSV of extensions to load.')
 177   (options, _) = parser.parse_args()
 178   if not options.extension_csv:
 179     parser.error('Must specify --extension-csv option.')
 180   return options.extension_csv
 181
 182 def _GetVersionFromCrx(crx_path):
 183   """Retrieves extension version from CRX archive.
 184
 185   Args:
 186     crx_path: path to CRX archive to extract version from.
 187   """
 188   with zipfile.ZipFile(crx_path, 'r') as crx_zip:
 189     manifest_contents = crx_zip.read('manifest.json')
 190     version = json.loads(manifest_contents)['version']
 191   return version
 192
 193 def _Base64Hash(file_path):
 194   return base64.b64encode(cloud_storage.CalculateHash(file_path))
 195
 196 def main():
 197   extension_csv = _GetCsvFromArgs()
 198   local_extensions_dir = tempfile.mkdtemp()
 199   try:
 200     _UpdateExtensionsInCloud(local_extensions_dir,
 201                              extension_csv, REMOTE_DIR)
 202   finally:
 203     shutil.rmtree(local_extensions_dir)
 204
 205 if __name__ == '__main__':
 206   main()
 207