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