hash.py

import os
import sys
import hashlib
import argparse

BUFFER_SIZE = 4 * 1024


class HashResult(object):
    def __init__(self, path, hasher):
        """Creates new HashResult

        :param path: the path of the file
        :type path: str

        :param hasher: the resulting hash object
        """
        self.path = path
        self.hasher = hasher


def main():
    args = get_args()

    for path in args.paths:
        hash_generator = calculate_hash(args.hash_method, path, args.recursive)

        column_len = 0
        i = 0
        for result in hash_generator:
            file_name = os.path.basename(result.path)
            cur_column_len = len(file_name)
            if cur_column_len > column_len:
                column_len = cur_column_len
            formatting = "{0:" + str(column_len) + "} {1}"

            print formatting.format(file_name, result.hasher.hexdigest())
            i += 1

        print "\n%i files were hashed." % i


def calculate_hash(method, path, is_recursive):
    """A generator that calculates the hash of the path using the hashing method.

    :param method: the hashing method to use
    :type method: str

    :param path: the path to hash
    :type path: str

    :param is_recursive: if recursive hashing is enabled
    :type is_recursive: bool

    :return: the generator that yields a :ref
    """
    try:
        # Try to append the working directory if the path is relative
        if not os.path.isfile(path) or os.path.isdir(path):
            path = os.path.join(os.getcwd(), path)

        # Check if the path is actually a folder
        if os.path.isdir(path):
            if not is_recursive:
                print "Info: Recursive scanning is disabled. Skipping directory '%s'.\n" \
                      "      Specify option '-r' to enable recursive scanning." % path
                return
            else:
                for p in os.listdir(path):
                    generator = calculate_hash(method, os.path.join(path, p), is_recursive)
                    for result in generator:
                        yield result
                return

        hasher = get_hasher(method)
        size = os.path.getsize(path)
        f = open(path, 'r')

        i = 0
        while i < size:
            data = f.read(BUFFER_SIZE)
            hasher.update(data)
            i += BUFFER_SIZE

        yield HashResult(path, hasher)
    except (IOError, OSError):
        print "Error: Unable to read '%s'." % path


def get_hasher(hash_method):
    try:
        return hashlib.new(hash_method)
    except ValueError:
        print "Error: Unsupported hash method '%s'." % hash_method
        sys.exit(-1)


def get_args():
    # hash_method = sys.argv[1]
    # files = sys.argv[2:]

    parser = argparse.ArgumentParser()
    parser.add_argument('hash_method', metavar='method',
                        help="the hashing method to use. Can be a value of: " +
                             str(hashlib.algorithms)[1:-1].replace("'", ""))
    parser.add_argument('paths', metavar='path', nargs='+',
                        help="path to run the hash on")
    parser.add_argument('-r', '--recursive', dest='recursive',
                        help="enables all sub-files and sub-directories to be hashed when a "
                             "directory is specified",
                        action="store_true")
    return parser.parse_args()

if __name__ == '__main__':
    main()