compute-file-hashes.py

#!/usr/bin/env python
import os
import hashlib
import argparse
from exitstatus import ExitStatus


# Customized exit handler, though we can't let it be called anywhere else.
def _exit(error_occurred=False, err=None):
    if error_occurred:
        print "\r\nProcess ended in error, details:"
        print _get_exception_string(err)
        exit(ExitStatus.failure)
    else:
        exit(ExitStatus.success)


# Custom built-in for formatting Exception output in a friendly-ish way.
def _get_exception_string(exception=None):
    # errtype=type(e).__name__, errmsg=str(e)
    # print "[%s] %s" % (errtype, errmsg)
    return "[%s] %s" % (type(exception).__name__, str(e))


# Return the size of the file in bytes, necessary for chunking determination
def _get_filesize(path):
    return os.path.getsize(path)


# Handler for hash computations, so we don't have to repeat the same 'hash it' code
# four times over
#
# noinspection PyUnusedLocal,PyShadowingNames
def _compute_hash(filepath, digest):
    try:
        # Init the hashlib object
        if digest == 'md5':
            filehash = hashlib.md5()
        elif digest == 'sha1':
            filehash = hashlib.sha1()
        elif digest == 'sha256':
            filehash = hashlib.sha256()
        elif digest == 'sha512':
            filehash = hashlib.sha512()
        else:
            raise TypeError(
                "Invalid digest type specified, accepted are: md5 , sha1, sha256, sha512")

        # Open the file specified in the filepath.
        with open(filepath, "rb") as filename:
            # Check file size, and chunk it if it's greater than 1MB
            if _get_filesize(filepath) > 1000000:
                for chunk in iter(lambda: filename.read(4096), b""):
                    # For each chunk, incorporate the chunk into the hash sum.
                    filehash.update(chunk)
            else:
                # If we don't need to chunk, run the hash sum directly on the file (<= 1MB only)
                filehash.update(filename.read())

        # Return the MD5 hash sum (hexadecimal) to the end user.
        return filehash.hexdigest()
    except Exception as e:
        raise RuntimeError(str(e))


# noinspection PyShadowingNames,PyProtectedMember
def _get_arguments():
    # Argument Parser for the overall function
    parser = argparse.ArgumentParser(
        description="Run hash sum functions on a file and return them (default: "
                    "run both MD5 and SHA1, if no hash sum flags provided).",
        add_help=True)

    # Reset the title for "Optional Arguments" (help builtin)
    # NOTE: May introduce undefined behavior in future, be careful!
    parser._optionals.title = "Optional Arguments"

    # Argument Group for Hash Sum Flags, which define what hashes to run when at least one
    # of these args is specified.
    sum_flags = parser.add_argument_group('Hash Sum Flags')
    sum_flags.add_argument('-md5', required=False, action="store_true",
                           default=False,
                           help="Generate an MD5 sum for the file.")
    sum_flags.add_argument('-sha1', required=False, action="store_true",
                           default=False,
                           help="Generate a SHA1 sum for the file.")
    sum_flags.add_argument('-sha256', required=False, action="store_true",
                           default=False,
                           help="Generate a SHA256 sum for the file.")
    sum_flags.add_argument('-sha512', required=False, action="store_true",
                           default=False,
                           help="Generate a SHA512 sum for the file.")

    # Argument Group for Required Arguments, i.e. the path to the file.
    required_args = parser.add_argument_group('Required Arguments')
    required_args.add_argument('filepath', help="Path to file to run sums on.",
                               metavar="filepath")

    # Actually parse the args, and return this to the system.
    return parser.parse_args()


def main():
    # Init the vars for sums, used for checks later.
    md5 = None
    sha1 = None
    sha256 = None
    sha512 = None

    try:
        # Provide some useful data to user about progress, and what file we're actually checking.
        print "FilePath: %s\r\n" % filepath
        if "md5" in SUMS_TO_RUN:
            print "Generating MD5 sum, this could take some time..."
            # Actually compute MD5
            md5 = _compute_hash(filepath, 'md5')
            # More useful progress data.
            print "MD5 hash calculated.\r\n"

        if "sha1" in SUMS_TO_RUN:
            print "Generating SHA1 sum, this could take some time..."
            # Actually compute SHA1
            sha1 = _compute_hash(filepath, 'sha1')
            # More useful progress data.
            print "SHA1 hash calculated.\r\n"

        if "sha256" in SUMS_TO_RUN:
            print "Generating sha256 sum, this could take some time..."
            # Actually compute sha256
            sha256 = _compute_hash(filepath, 'sha256')
            # More useful progress data.
            print "SHA256 hash calculated.\r\n"

        if "sha512" in SUMS_TO_RUN:
            print "Generating SHA512 sum, this could take some time..."
            # Actually compute sha512
            sha512 = _compute_hash(filepath, 'sha512')
            # More useful progress data.
            print "SHA512 hash calculated.\r\n"

        print "Done, output below.\r\n\r\n"
        print "FilePath: %s\r\n" % filepath
        print "Requested Hash Sums:"
        if md5:
            print "    MD5 sum: %s" % str(md5)
        if sha1:
            print "   SHA1 sum: %s" % str(sha1)
        if sha256:
            print " SHA256 sum: %s" % str(sha256)
        if sha512:
            print " SHA512 sum: %s" % str(sha512)

    except Exception as err:
        # If the process dies off for any reason, we need to catch the exception and state it back
        # to the user, then die off gracefully.
        print "Process died with error: %s" % str(err)
        _exit(error_occurred=True, err=err)

    _exit()


# We have a few things we need to run *before* we get to main().
# Namely, argument parsing, cleaning the file paths slightly, and
# also checking if the file path exists for checking.
if __name__ == "__main__":
    # noinspection PyBroadException
    try:
        # Init "SUMS_TO_RUN", which will contain what hash sums we should run.
        SUMS_TO_RUN = []

        # First, we need to get the arguments.
        args = _get_arguments()

        # Store the filepath argument.
        filepath = args.filepath

        # First check if the filepath contains any slashes of any kind, and if not, rewrite path
        # to cwd.  Also triggers if path is `./blah`.
        if ('/' not in filepath and '\\' not in filepath) or ('./' in filepath):
            if '/' in os.getcwd():
                filepath = os.getcwd() + '/' + filepath
            elif '\\' in os.getcwd():
                filepath = os.getcwd() + '\\' + filepath

        # Now we check if the path provided exists
        if not os.path.exists(filepath):
            raise IOError(
                "The specified file path does not exist. Provided path: %s" % filepath)

        # PreProcessing: Determine from args which functions to run.
        if not (args.md5 or args.sha1 or args.sha256 or args.sha512):
            SUMS_TO_RUN = ['md5', 'sha1']
        else:
            if args.md5:
                SUMS_TO_RUN.append('md5')
            if args.sha1:
                SUMS_TO_RUN.append('sha1')
            if args.sha256:
                SUMS_TO_RUN.append('sha256')
            if args.sha512:
                SUMS_TO_RUN.append('sha512')

        # Now, call main() which will execute the actual hash functions.
        main()

    except Exception as e:
        _exit(error_occurred=True, err=e)

    else:
        _exit()