1import sys
2
3def calc_checksum(data, old):
4	a = 0
5	c = 0
6	for b in data:
7		a = a + b + c
8		c = 0
9		if a > 255:
10			a = a - 256
11			c = 1
12	if not old:
13		a += c
14	return a
15
16def usage():
17	print "Usage: python checksum.py {--old|--new} <file.bin> [<checksum> <offset>]"
18	exit(1)
19
20# main
21
22if len(sys.argv) != 3 and len(sys.argv) != 5:
23	usage()
24
25if sys.argv[1] == "--old":
26	old = True
27elif sys.argv[1] == "--new":
28	old = False
29else:
30	usage()
31
32filename = sys.argv[2]
33
34data = bytearray(open(filename, 'rb').read())
35
36checksum = calc_checksum(data, old)
37
38if len(sys.argv) != 5:
39	print filename + ": " + hex(checksum)
40	exit(0)
41
42desired_checksum = int(sys.argv[3], 16)
43offset = int(sys.argv[4], 16)
44
45if checksum < desired_checksum:
46	data[offset] = desired_checksum - checksum
47else:
48	data[offset] = 0xff - (checksum - desired_checksum)
49
50file = open(filename, "wb")
51file.write(data)
52