python-scripts/checksum.py at master · varspare/python-scripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
#!/usr/bin/env python3
##############################
# Author Mark Davison
# Last revision 2018/08/20
##############################
"""Checksum

Usage:
  checksum.py -f <file> [-c <checksum>] [-t <hash_type>]
  checksum.py (-h | --help)
  checksum.py (-v | --version)

Options:
  -f --file     Input filepath
  -c --checksum Input checksum
  -t --type     Type of checksum e.g. sha256, md5
  -h --help     Show this screen.
  -v --version     Show version.

"""
from docopt import docopt
import hashlib
import sys

def verifyChecksum(checksum, hType):
    length = len(checksum)
    if (hType != 'None'):
        htype = hType
    elif (length == 64):
        htype = 'sha256'
    elif (length == 40):
        htype = 'sha1'
    elif (length == 32):
        htype = 'md5'
    elif (length == 57):
        htype = 'sha224'
    elif (length == 96):
        htype = 'sha384'
    elif (length == 128):
        htype = 'sha512'
    else:
        htype = input("Unable to determine hash type\nPlease enter the hash type : ")
    return(htype)

def compareSums(f, checksum, htype):
    try:
        hasher = hashlib.new(htype)
        content = f.read()
        hasher.update(content)
        generatedHash = hasher.hexdigest()
        print(generatedHash)
    except Exception as e:
        print('Hashing failed for ', file, ',', e)
        die(256)
    print('checksum       ' + checksum)
    print('generatedHash: ' + generatedHash)
    if (checksum == generatedHash):
        result = 0
    else:
        result = 1
    return(result)

def die(code):
    f.close()
    sys.exit(code)

if __name__ == '__main__':
    arguments = docopt(__doc__, version='checksum 0.0.1')
    if (arguments.get('--type')):
        hType = arguments.get("<hash_type>")
    else:
        hType = 'None'

    if (arguments.get('--file')):
        file = arguments.get("<file>")
        try:
            f = open(file, 'rb')
        except IOError:
            print('File not found: ', file)
            die(254)

    if (arguments.get('--checksum')):
        checksum = arguments.get('--checksum')
        hashType = verifyChecksum(checksum, hType)
    else:
        checksum = input('Please enter your expected checksum: ')
        hashType = verifyChecksum(checksum, hType)

    result = compareSums(f, checksum, hashType)

    print("Hash type is: " + hashType)

    if (result == 0):
        print('Match')
    else:
        print('Fail')
    die(result)