michael@0: #! /usr/bin/env python
michael@0: # This Source Code Form is subject to the terms of the Mozilla Public
michael@0: # License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0: # file, You can obtain one at http://mozilla.org/MPL/2.0/.
michael@0: 
michael@0: #
michael@0: # Script name:   talos_from_code.py
michael@0: # Purpose:       Read from a talos.json file the different files to download for a talos job
michael@0: # Author(s):     Zambrano Gasparnian, Armen <armenzg@mozilla.com>
michael@0: # Target:        Python 2.5
michael@0: #
michael@0: from optparse import OptionParser
michael@0: import json
michael@0: import re
michael@0: import urllib2
michael@0: import urlparse
michael@0: import sys
michael@0: import os
michael@0: 
michael@0: def main():
michael@0:     '''
michael@0:     This script downloads a talos.json file which indicates which files to download
michael@0:     for a talos job.
michael@0:     See a talos.json file for a better understand:
michael@0:     http://hg.mozilla.org/mozilla-central/raw-file/default/testing/talos/talos.json
michael@0:     '''
michael@0:     parser = OptionParser()
michael@0:     parser.add_option("--talos-json-url", dest="talos_json_url", type="string",
michael@0:                       help="It indicates from where to download the talos.json file.")
michael@0:     (options, args) = parser.parse_args()
michael@0: 
michael@0:     # 1) check that the url was passed
michael@0:     if options.talos_json_url == None:
michael@0:         print "You need to specify --talos-json-url."
michael@0:         sys.exit(1)
michael@0: 
michael@0:     # 2) try to download the talos.json file
michael@0:     try:
michael@0:         jsonFilename = download_file(options.talos_json_url)
michael@0:     except Exception, e:
michael@0:         print "ERROR: We tried to download the talos.json file but something failed."
michael@0:         print "ERROR: %s" % str(e)
michael@0:         sys.exit(1)
michael@0: 
michael@0:     # 3) download the necessary files
michael@0:     print "INFO: talos.json URL: %s" % options.talos_json_url
michael@0:     try:
michael@0:         key = 'talos.zip'
michael@0:         entity = get_value(jsonFilename, key)
michael@0:         if passesRestrictions(options.talos_json_url, entity["url"]):
michael@0:             # the key is at the same time the filename e.g. talos.zip
michael@0:             print "INFO: Downloading %s as %s" % (entity["url"], os.path.join(entity["path"], key))
michael@0:             download_file(entity["url"], entity["path"], key)
michael@0:         else:
michael@0:             print "ERROR: You have tried to download a file " + \
michael@0:                   "from: %s " % fileUrl + \
michael@0:                   "which is a location different than http://talos-bundles.pvt.build.mozilla.org/"
michael@0:             print "ERROR: This is only allowed for the certain branches."
michael@0:             sys.exit(1)
michael@0:     except Exception, e:
michael@0:         print "ERROR: %s" % str(e)
michael@0:         sys.exit(1)
michael@0: 
michael@0: def passesRestrictions(talosJsonUrl, fileUrl):
michael@0:     '''
michael@0:     Only certain branches are exempted from having to host their downloadable files
michael@0:     in talos-bundles.pvt.build.mozilla.org
michael@0:     '''
michael@0:     if talosJsonUrl.startswith("http://hg.mozilla.org/try/") == True or \
michael@0:        talosJsonUrl.startswith("https://hg.mozilla.org/try/") == True or \
michael@0:        talosJsonUrl.startswith("http://hg.mozilla.org/projects/pine/") == True or \
michael@0:        talosJsonUrl.startswith("https://hg.mozilla.org/projects/pine/") == True:
michael@0:         return True
michael@0:     else:
michael@0:         p = re.compile('^http://talos-bundles.pvt.build.mozilla.org/')
michael@0:         m = p.match(fileUrl)
michael@0:         if m == None:
michael@0:             return False
michael@0:         return True
michael@0: 
michael@0: def get_filename_from_url(url):
michael@0:     '''
michael@0:     This returns the filename of the file we're trying to download
michael@0:     '''
michael@0:     parsed = urlparse.urlsplit(url.rstrip('/'))
michael@0:     if parsed.path != '':
michael@0:         return parsed.path.rsplit('/', 1)[-1]
michael@0:     else:
michael@0:         print "ERROR: We were trying to download a file from %s " + \
michael@0:               "but the URL seems to be incorrect."
michael@0:         sys.exit(1)
michael@0: 
michael@0: def download_file(url, path="", saveAs=None):
michael@0:     '''
michael@0:     It downloads a file from URL to the indicated path
michael@0:     '''
michael@0:     req = urllib2.Request(url)
michael@0:     f = urllib2.urlopen(req)
michael@0:     if path != "" and not os.path.isdir(path):
michael@0:         try:
michael@0:             os.makedirs(path)
michael@0:             print "INFO: directory %s created" % path
michael@0:         except Exception, e:
michael@0:             print "ERROR: %s" % str(e)
michael@0:             sys.exit(1)
michael@0:     filename = saveAs if saveAs else get_filename_from_url(url)
michael@0:     local_file = open(os.path.join(path, filename), 'wb')
michael@0:     local_file.write(f.read())
michael@0:     local_file.close()
michael@0:     return filename
michael@0: 
michael@0: def get_value(json_filename, key):
michael@0:     '''
michael@0:     It loads up a JSON file and returns the value for the given string
michael@0:     '''
michael@0:     f = open(json_filename, 'r')
michael@0:     return json.load(f)[key]
michael@0: 
michael@0: if __name__ == '__main__':
michael@0:     main()