build/checksums.py

   1 #!/usr/bin/python
   2 # This Source Code Form is subject to the terms of the Mozilla Public
   3 # License, v. 2.0. If a copy of the MPL was not distributed with this
   4 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
   5
   6 from __future__ import with_statement
   7
   8 from optparse import OptionParser
   9 import hashlib
  10 import logging
  11 import os
  12
  13 logger = logging.getLogger('checksums.py')
  14
  15
  16 def digest_file(filename, digest, chunk_size=131072):
  17     '''Produce a checksum for the file specified by 'filename'.  'filename'
  18     is a string path to a file that is opened and read in this function.  The
  19     checksum algorithm is specified by 'digest' and is a valid OpenSSL
  20     algorithm.  If the digest used is not valid or Python's hashlib doesn't
  21     work, the None object will be returned instead.  The size of blocks
  22     that this function will read from the file object it opens based on
  23     'filename' can be specified by 'chunk_size', which defaults to 1K'''
  24     assert not os.path.isdir(filename), 'this function only works with files'
  25
  26     logger.debug('Creating new %s object' % digest)
  27     h = hashlib.new(digest)
  28     with open(filename, 'rb') as f:
  29         while True:
  30             data = f.read(chunk_size)
  31             if not data:
  32                 logger.debug('Finished reading in file')
  33                 break
  34             h.update(data)
  35     hash = h.hexdigest()
  36     logger.debug('Hash for %s is %s' % (filename, hash))
  37     return hash
  38
  39
  40 def process_files(dirs, output_filename, digests):
  41     '''This function takes a list of directory names, 'drs'. It will then
  42     compute the checksum for each of the files in these by by opening the files.
  43     Once each file is read and its checksum is computed, this function
  44     will write the information to the file specified by 'output_filename'.
  45     The path written in the output file will have anything specified by 'strip'
  46     removed from the path.  The output file is closed before returning nothing
  47     The algorithm to compute checksums with can be specified by 'digests'
  48     and needs to be a list of valid OpenSSL algorithms.
  49
  50     The output file is written in the format:
  51         <hash> <algorithm> <filesize> <filepath>
  52     Example:
  53         d1fa09a<snip>e4220 sha1 14250744 firefox-4.0b6pre.en-US.mac64.dmg
  54     '''
  55
  56     if os.path.exists(output_filename):
  57         logger.debug('Overwriting existing checksums file "%s"' %
  58                      output_filename)
  59     else:
  60         logger.debug('Creating a new checksums file "%s"' % output_filename)
  61     with open(output_filename, 'w+') as output:
  62         for d in dirs:
  63             for root, dirs, files in os.walk(d):
  64                 for f in files:
  65                     full = os.path.join(root, f)
  66                     rel = os.path.relpath(full, d)
  67
  68                     for digest in digests:
  69                         hash = digest_file(full, digest)
  70
  71                         output.write('%s %s %s %s\n' % (
  72                             hash, digest, os.path.getsize(full), rel))
  73
  74
  75 def setup_logging(level=logging.DEBUG):
  76     '''This function sets up the logging module using a speficiable logging
  77     module logging level.  The default log level is DEBUG.
  78
  79     The output is in the format:
  80         <level> - <message>
  81     Example:
  82         DEBUG - Finished reading in file
  83 '''
  84
  85     logger = logging.getLogger('checksums.py')
  86     logger.setLevel(logging.DEBUG)
  87     handler = logging.StreamHandler()
  88     handler.setLevel(level)
  89     formatter = logging.Formatter("%(levelname)s - %(message)s")
  90     handler.setFormatter(formatter)
  91     logger.addHandler(handler)
  92
  93
  94 def main():
  95     '''This is a main function that parses arguments, sets up logging
  96     and generates a checksum file'''
  97     # Parse command line arguments
  98     parser = OptionParser()
  99     parser.add_option('-d', '--digest', help='checksum algorithm to use',
 100                       action='append', dest='digests')
 101     parser.add_option('-o', '--output', help='output file to use',
 102                       action='store', dest='outfile', default='checksums')
 103     parser.add_option('-v', '--verbose',
 104                       help='Be noisy (takes precedence over quiet)',
 105                       action='store_true', dest='verbose', default=False)
 106     parser.add_option('-q', '--quiet', help='Be quiet', action='store_true',
 107                       dest='quiet', default=False)
 108
 109     options, args = parser.parse_args()
 110
 111     # Figure out which logging level to use
 112     if options.verbose:
 113         loglevel = logging.DEBUG
 114     elif options.quiet:
 115         loglevel = logging.ERROR
 116     else:
 117         loglevel = logging.INFO
 118
 119     # Set up logging
 120     setup_logging(loglevel)
 121
 122     # Validate the digest type to use
 123     if not options.digests:
 124         options.digests = ['sha1']
 125
 126     for i in args:
 127         if not os.path.isdir(i):
 128             logger.error('%s is not a directory' % i)
 129             exit(1)
 130
 131     process_files(args, options.outfile, options.digests)
 132
 133
 134 if __name__ == '__main__':
 135     main()