michael@0: #!/usr/bin/python
michael@0: #
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: import sys
michael@0: import os
michael@0: import hashlib
michael@0: import json
michael@0: import re
michael@0: import errno
michael@0: from argparse import ArgumentParser
michael@0: 
michael@0: def getFileHashAndSize(filename):
michael@0:     sha512Hash = 'UNKNOWN'
michael@0:     size = 'UNKNOWN'
michael@0: 
michael@0:     try:
michael@0:         # open in binary mode to make sure we get consistent results
michael@0:         # across all platforms
michael@0:         f = open(filename, "rb")
michael@0:         shaObj = hashlib.sha512(f.read())
michael@0:         sha512Hash = shaObj.hexdigest()
michael@0: 
michael@0:         size = os.path.getsize(filename)
michael@0:     except:
michael@0:         pass
michael@0: 
michael@0:     return (sha512Hash, size)
michael@0: 
michael@0: def getMarProperties(filename, partial=False):
michael@0:     if not os.path.exists(filename):
michael@0:         return {}
michael@0:     (mar_hash, mar_size) = getFileHashAndSize(filename)
michael@0:     martype = 'partial' if partial else 'complete'
michael@0:     return {
michael@0:         '%sMarFilename' % martype: os.path.basename(filename),
michael@0:         '%sMarSize' % martype: mar_size,
michael@0:         '%sMarHash' % martype: mar_hash,
michael@0:     }
michael@0: 
michael@0: def getUrlProperties(filename):
michael@0:     # let's create a switch case using name-spaces/dict
michael@0:     # rather than a long if/else with duplicate code
michael@0:     property_conditions = [
michael@0:         # key: property name, value: condition
michael@0:         ('symbolsUrl', lambda m: m.endswith('crashreporter-symbols.zip') or
michael@0:                        m.endswith('crashreporter-symbols-full.zip')),
michael@0:         ('testsUrl', lambda m: m.endswith(('tests.tar.bz2', 'tests.zip'))),
michael@0:         ('unsignedApkUrl', lambda m: m.endswith('apk') and
michael@0:                            'unsigned-unaligned' in m),
michael@0:         ('robocopApkUrl', lambda m: m.endswith('apk') and 'robocop' in m),
michael@0:         ('jsshellUrl', lambda m: 'jsshell-' in m and m.endswith('.zip')),
michael@0:         ('completeMarUrl', lambda m: m.endswith('.complete.mar')),
michael@0:         ('partialMarUrl', lambda m: m.endswith('.mar') and '.partial.' in m),
michael@0:         # packageUrl must be last!
michael@0:         ('packageUrl', lambda m: True),
michael@0:     ]
michael@0:     url_re = re.compile(r'''^(https?://.*?\.(?:tar\.bz2|dmg|zip|apk|rpm|mar|tar\.gz))$''')
michael@0:     properties = {}
michael@0: 
michael@0:     try:
michael@0:         with open(filename) as f:
michael@0:             for line in f:
michael@0:                 m = url_re.match(line)
michael@0:                 if m:
michael@0:                     m = m.group(1)
michael@0:                     for prop, condition in property_conditions:
michael@0:                         if condition(m):
michael@0:                             properties.update({prop: m})
michael@0:                             break
michael@0:     except IOError as e:
michael@0:         if e.errno != errno.ENOENT:
michael@0:             raise
michael@0:         properties = {prop: 'UNKNOWN' for prop, condition in property_conditions}
michael@0:     return properties
michael@0: 
michael@0: def getPartialInfo(props):
michael@0:     return [{
michael@0:         "from_buildid": props.get("previous_buildid"),
michael@0:         "size": props.get("partialMarSize"),
michael@0:         "hash": props.get("partialMarHash"),
michael@0:         "url": props.get("partialMarUrl"),
michael@0:     }]
michael@0: 
michael@0: if __name__ == '__main__':
michael@0:     parser = ArgumentParser(description='Generate mach_build_properties.json for automation builds.')
michael@0:     parser.add_argument("--complete-mar-file", required=True,
michael@0:                         action="store", dest="complete_mar_file",
michael@0:                         help="Path to the complete MAR file, relative to the objdir.")
michael@0:     parser.add_argument("--partial-mar-file", required=False,
michael@0:                         action="store", dest="partial_mar_file",
michael@0:                         help="Path to the partial MAR file, relative to the objdir.")
michael@0:     parser.add_argument("--upload-output", required=True,
michael@0:                         action="store", dest="upload_output",
michael@0:                         help="Path to the text output of 'make upload'")
michael@0:     args = parser.parse_args()
michael@0: 
michael@0:     json_data = getMarProperties(args.complete_mar_file)
michael@0:     json_data.update(getUrlProperties(args.upload_output))
michael@0:     if args.partial_mar_file:
michael@0:         json_data.update(getMarProperties(args.partial_mar_file, partial=True))
michael@0: 
michael@0:         # Pull the previous buildid from the partial mar filename.
michael@0:         res = re.match(r'.*\.([0-9]+)-[0-9]+.mar', args.partial_mar_file)
michael@0:         if res:
michael@0:             json_data['previous_buildid'] = res.group(1)
michael@0: 
michael@0:             # Set partialInfo to be a collection of the partial mar properties
michael@0:             # useful for balrog.
michael@0:             json_data['partialInfo'] = getPartialInfo(json_data)
michael@0: 
michael@0:     with open('mach_build_properties.json', 'w') as outfile:
michael@0:         json.dump(json_data, outfile, indent=4)