michael@0: #!/usr/bin/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: from __future__ import with_statement
michael@0: 
michael@0: from optparse import OptionParser
michael@0: import logging
michael@0: import os
michael@0: try:
michael@0:     import hashlib
michael@0: except:
michael@0:     hashlib = None
michael@0: 
michael@0: def digest_file(filename, digest, chunk_size=1024):
michael@0:     '''Produce a checksum for the file specified by 'filename'.  'filename'
michael@0:     is a string path to a file that is opened and read in this function.  The
michael@0:     checksum algorithm is specified by 'digest' and is a valid OpenSSL
michael@0:     algorithm.  If the digest used is not valid or Python's hashlib doesn't
michael@0:     work, the None object will be returned instead.  The size of blocks
michael@0:     that this function will read from the file object it opens based on
michael@0:     'filename' can be specified by 'chunk_size', which defaults to 1K'''
michael@0:     assert not os.path.isdir(filename), 'this function only works with files'
michael@0:     logger = logging.getLogger('checksums.py')
michael@0:     if hashlib is not None:
michael@0:         logger.debug('Creating new %s object' % digest)
michael@0:         h = hashlib.new(digest)
michael@0:         with open(filename, 'rb') as f:
michael@0:             while True:
michael@0:                 data = f.read(chunk_size)
michael@0:                 if not data:
michael@0:                     logger.debug('Finished reading in file')
michael@0:                     break
michael@0:                 h.update(data)
michael@0:         hash = h.hexdigest()
michael@0:         logger.debug('Hash for %s is %s' % (filename, hash))
michael@0:         return hash
michael@0:     else:
michael@0:         # In this case we could subprocess.Popen and .communicate with
michael@0:         # sha1sum or md5sum
michael@0:         logger.warn('The python module for hashlib is missing!')
michael@0:         return None
michael@0: 
michael@0: 
michael@0: def process_files(files, output_filename, digests, strip):
michael@0:     '''This function takes a list of file names, 'files'.  It will then
michael@0:     compute the checksum for each of the files by opening the files.
michael@0:     Once each file is read and its checksum is computed, this function
michael@0:     will write the information to the file specified by 'output_filename'.
michael@0:     The path written in the output file will have anything specified by 'strip'
michael@0:     removed from the path.  The output file is closed before returning nothing
michael@0:     The algorithm to compute checksums with can be specified by 'digests' 
michael@0:     and needs to be a list of valid OpenSSL algorithms.
michael@0: 
michael@0:     The output file is written in the format:
michael@0:         <hash> <algorithm> <filesize> <filepath>
michael@0:     Example:
michael@0:         d1fa09a<snip>e4220 sha1 14250744 firefox-4.0b6pre.en-US.mac64.dmg
michael@0:     '''
michael@0: 
michael@0:     logger = logging.getLogger('checksums.py')
michael@0:     if os.path.exists(output_filename):
michael@0:         logger.debug('Overwriting existing checksums file "%s"' %
michael@0:                      output_filename)
michael@0:     else:
michael@0:         logger.debug('Creating a new checksums file "%s"' % output_filename)
michael@0:     with open(output_filename, 'w+') as output:
michael@0:         for file in files:
michael@0:             if os.path.isdir(file):
michael@0:                 logger.warn('%s is a directory, skipping' % file)
michael@0:             else:
michael@0:                 for digest in digests:
michael@0:                     hash = digest_file(file, digest)
michael@0:                     if hash is None:
michael@0:                         logger.warn('Unable to generate a hash for %s. ' +
michael@0:                                     'Skipping.' % file)
michael@0:                         continue
michael@0:                     if file.startswith(strip):
michael@0:                         short_file = file[len(strip):]
michael@0:                         short_file = short_file.lstrip('/')
michael@0:                     else:
michael@0:                         short_file = file
michael@0:                     print >>output, '%s %s %s %s' % (hash, digest,
michael@0:                                                      os.path.getsize(file),
michael@0:                                                      short_file)
michael@0: 
michael@0: def setup_logging(level=logging.DEBUG):
michael@0:     '''This function sets up the logging module using a speficiable logging
michael@0:     module logging level.  The default log level is DEBUG.
michael@0: 
michael@0:     The output is in the format:
michael@0:         <level> - <message>
michael@0:     Example:
michael@0:         DEBUG - Finished reading in file
michael@0: '''
michael@0: 
michael@0:     logger = logging.getLogger('checksums.py')
michael@0:     logger.setLevel(logging.DEBUG)
michael@0:     handler = logging.StreamHandler()
michael@0:     handler.setLevel(level)
michael@0:     formatter = logging.Formatter("%(levelname)s - %(message)s")
michael@0:     handler.setFormatter(formatter)
michael@0:     logger.addHandler(handler)
michael@0: 
michael@0: def main():
michael@0:     '''This is a main function that parses arguments, sets up logging
michael@0:     and generates a checksum file'''
michael@0:     # Parse command line arguments
michael@0:     parser = OptionParser()
michael@0:     parser.add_option('-d', '--digest', help='checksum algorithm to use',
michael@0:                       action='append', dest='digests')
michael@0:     parser.add_option('-o', '--output', help='output file to use',
michael@0:                       action='store', dest='outfile', default='checksums')
michael@0:     parser.add_option('-v', '--verbose',
michael@0:                       help='Be noisy (takes precedence over quiet)',
michael@0:                       action='store_true', dest='verbose', default=False)
michael@0:     parser.add_option('-q', '--quiet', help='Be quiet', action='store_true',
michael@0:                       dest='quiet', default=False)
michael@0:     parser.add_option('-s', '--strip',
michael@0:                       help='strip this path from the filenames',
michael@0:                       dest='strip', default=os.getcwd())
michael@0:     options, args = parser.parse_args()
michael@0: 
michael@0:     #Figure out which logging level to use
michael@0:     if options.verbose:
michael@0:         loglevel = logging.DEBUG
michael@0:     elif options.quiet:
michael@0:         loglevel = logging.ERROR
michael@0:     else:
michael@0:         loglevel = logging.INFO
michael@0: 
michael@0:     #Set up logging
michael@0:     setup_logging(loglevel)
michael@0:     logger = logging.getLogger('checksums.py')
michael@0: 
michael@0:     # Validate the digest type to use
michael@0:     if not options.digests:
michael@0:         options.digests = ['sha1']
michael@0:     try:
michael@0:         for digest in options.digests:
michael@0:             hashlib.new(digest)
michael@0:     except ValueError, ve:
michael@0:         logger.error('Could not create a "%s" hash object (%s)' %
michael@0:                      (digest, ve.args[0]))
michael@0:         exit(1)
michael@0: 
michael@0:     # Validate the files to checksum
michael@0:     files = []
michael@0:     for i in args:
michael@0:         if os.path.exists(i):
michael@0:             files.append(i)
michael@0:         else:
michael@0:             logger.info('File "%s" was not found on the filesystem' % i)
michael@0:     process_files(files, options.outfile, options.digests, options.strip)
michael@0: 
michael@0: if __name__ == '__main__':
michael@0:     main()