pipe-checksum

#!/usr/bin/env python3

import argparse
import hashlib
import sys

def main():
    parser = argparse.ArgumentParser(
        description='Perform a checksum on stdin to validate `curl | bash` installs')
    parser.add_argument('-a', '--algorithm', dest='algorithm', choices=[ '1', '256', '512', 'md5' ],
                        required=False, default='1', help='The checksum algorithm to apply')
    parser.add_argument('checksum', type=str, nargs=1, help='The expected checksum of stdin')

    args = parser.parse_args()

    if args.algorithm == '1':
        hasher = hashlib.sha1()
    elif args.algorithm == '256':
        hasher = hashlib.sha256()
    elif args.algorithm == '512':
        hasher = hashlib.sha512()
    elif args.algorithm == 'md5':
        hasher = hashlib.md5()
    else:
        raise Exception("Unexpected algorithm: %s" % args.algorithm) # Parser should prevent this

    data = sys.stdin.buffer.read()
    hasher.update(data)
    digest = hasher.hexdigest()
    if digest == args.checksum[0]:
        sys.stdout.buffer.write(data)
    else:
        print('stdin checksum did not match! Got: %s' % digest)
        return -1

if __name__ == '__main__':
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        sys.exit(130)