checksum.py

import hashlib
import optparse
from os import walk
from os.path import dirname, isfile, join, relpath,basename

HASH_BLOCK_SIZE = 65536


def hash_file(fname, dig):
    if dig == "sha256":
        algorithm = hashlib.sha256()
    if dig == "sha512":
        algorithm = hashlib.sha3_512()
    if dig == "sha1":
        algorithm = hashlib.sha1()
    if dig == "md5":
        algorithm = hashlib.md5()

    with open(fname, "rb") as f:
        for chunk in iter(lambda: f.read(HASH_BLOCK_SIZE), b""):
            algorithm.update(chunk)
    return algorithm.hexdigest()


def verify_checksums(path):
    checksum_filename = path
    dirpath = dirname(path)
    bname = basename(path)
    dig = bname.split('.')[1]

    errors = {}
    with open(checksum_filename, mode="r") as checksum_list:
        for line in checksum_list:
            expected_hash, filename = line.strip().split(" ", 1)
            if filename.startswith("*"):
                filename = filename[1:]
            actual_hash = hash_file(join(dirpath, filename), dig)

            if expected_hash != actual_hash:
                errors[filename] = 'FAIL'
                print(f'{filename} FAIL')

    return errors


def calc_file_hash(filepath, dig):
    with open(f'checksum.{dig}', 'w') as checksumfile:
        checksumfile.write(f"{hash_file(filepath,dig)} *{filepath}")
        checksumfile.write('\n')
    return filepath


def calc_dir_hash(path, dig):
    rootpath = path
    with open(join(rootpath, f'checksum.{dig}'), 'w') as checksumfile:
        for (dirpath, dirnames, filenames) in walk(rootpath):
            for file in filenames:
                if file.endswith(f'.{dig}'):
                    continue
                rel_path = relpath(dirpath, rootpath)
                try:
                    checksumfile.write(f"{hash_file(join(dirpath, file),dig)} *{join(rel_path, file)}")
                    #print(f"{hash_file(join(dirpath, file),dig)} *{join(rel_path, file)}")
                    checksumfile.write('\n')
                except:
                    continue


def main():
    prog = optparse.OptionParser('A simple tool for fixity control')
    prog.add_option('-a', '--algorithm', default="md5",
                    choices=("md5", "sha1", "sha256", "sha512"),
                    help="Specifiy hashing algorithm md5, sha1, sha256, sha512 (default md5)")
    prog.add_option('-v', '--verify', action='store_true', help="Takes a checksum file and verifies calculated against "
                                                                "against expected checksum ")
    (opts, args) = prog.parse_args()

    if opts.verify:
        if len(args) != 1:
            prog.error('you must supply checksum file path')
        verify = verify_checksums(args[0])

        if len(verify) == 0:
            print('PASSED')

    else:
        if len(args) != 1:
            prog.error('you must supply input path')

        if isfile(args[0]):
            calc_file_hash(args[0], opts.algorithm)

        else:
            calc_dir_hash(args[0], opts.algorithm)


if __name__ == '__main__':
    main()