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"' % output_filename)
  58     else:
  59         logger.debug('Creating a new checksums file "%s"' % output_filename)
  60     with open(output_filename, "w+") as output:
  61         for d in dirs:
  62             for root, dirs, files in os.walk(d):
  63                 for f in files:
  64                     full = os.path.join(root, f)
  65                     rel = os.path.relpath(full, d)
  66
  67                     for digest in digests:
  68                         hash = digest_file(full, digest)
  69
  70                         output.write(
  71                             "%s %s %s %s\n" % (hash, digest, os.path.getsize(full), rel)
  72                         )
  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     logger = logging.getLogger("checksums.py")
  85     logger.setLevel(logging.DEBUG)
  86     handler = logging.StreamHandler()
  87     handler.setLevel(level)
  88     formatter = logging.Formatter("%(levelname)s - %(message)s")
  89     handler.setFormatter(formatter)
  90     logger.addHandler(handler)
  91
  92
  93 def main():
  94     """This is a main function that parses arguments, sets up logging
  95     and generates a checksum file"""
  96     # Parse command line arguments
  97     parser = OptionParser()
  98     parser.add_option(
  99         "-d",
 100         "--digest",
 101         help="checksum algorithm to use",
 102         action="append",
 103         dest="digests",
 104     )
 105     parser.add_option(
 106         "-o",
 107         "--output",
 108         help="output file to use",
 109         action="store",
 110         dest="outfile",
 111         default="checksums",
 112     )
 113     parser.add_option(
 114         "-v",
 115         "--verbose",
 116         help="Be noisy (takes precedence over quiet)",
 117         action="store_true",
 118         dest="verbose",
 119         default=False,
 120     )
 121     parser.add_option(
 122         "-q",
 123         "--quiet",
 124         help="Be quiet",
 125         action="store_true",
 126         dest="quiet",
 127         default=False,
 128     )
 129
 130     options, args = parser.parse_args()
 131
 132     # Figure out which logging level to use
 133     if options.verbose:
 134         loglevel = logging.DEBUG
 135     elif options.quiet:
 136         loglevel = logging.ERROR
 137     else:
 138         loglevel = logging.INFO
 139
 140     # Set up logging
 141     setup_logging(loglevel)
 142
 143     # Validate the digest type to use
 144     if not options.digests:
 145         options.digests = ["sha1"]
 146
 147     for i in args:
 148         if not os.path.isdir(i):
 149             logger.error("%s is not a directory" % i)
 150             exit(1)
 151
 152     process_files(args, options.outfile, options.digests)
 153
 154
 155 if __name__ == "__main__":
 156     main()