Checksum improvements

Makefile

@@ -1,14 +1,17 @@
+%.bin: %.s
+	ca65 -g "$<"
+
 all:
 	ca65 -g basic/basic.s
 	ca65 -g kernal/kernal.s
-	ld65 -C rom.cfg -o rom.bin basic/basic.o kernal/kernal.o -Ln rom.txt
+	ld65 -C rom.cfg -o /dev/null basic/basic.o kernal/kernal.o -Ln rom.txt
 	# if it's the unchanged 901226-01 image, use old checksum algorithm
 	$$SHELL -c "if [ $$(crc32 basic.bin) == cfdebff8 ]; then \
-		python checksum.py --old basic.bin 0xa0 0x1f52; \
+		python3 checksum.py -y --old basic.bin 0xa0 0x1f52; \
 	else \
-		python checksum.py --new basic.bin 0xa0 0x1f52; \
+		python3 checksum.py -y --new basic.bin 0xa0 0x1f52; \
 	fi"
-	python checksum.py --new kernal.bin 0xe0 0x4ac
+	python3 checksum.py -y --new kernal.bin 0xe0 0x4ac
 
 clean:
 	rm -f basic/basic.o kernal/kernal.o basic.bin kernal.bin rom.bin

checksum.py

@@ -1,6 +1,7 @@
+#!/usr/bin/env python3
 from __future__ import print_function
 
-import sys
+import sys, argparse
 
 def calc_checksum(data, old):
 	a = 0
@@ -15,39 +16,96 @@ def calc_checksum(data, old):
 		a += c
 	return a
 
-def usage():
-	print("Usage: python checksum.py {--old|--new} <file.bin> [<checksum> <offset>]")
-	exit(1)
+def to_base(num:int,base:int):
+	if base > 36: raise NotImplementedError("only bases up to 36")
+	digits="0123456789abcdefghijklmnopqrstuvwxyz"
+	ret = ""
+	remainder = int(abs(num))
+	if remainder==0: return "0"
+	while remainder>0:
+		ret = digits[remainder % base] + ret
+		remainder=remainder // base	
+	return ret
+
+def confirm(msg):	
+	r=input(f"{msg} (Y/N)? ").strip().lower()
+	if r != "y": return False
+	return True
 
 # main
 
-if len(sys.argv) != 3 and len(sys.argv) != 5:
-	usage()
+if __name__ == "__main__":
+	parser = argparse.ArgumentParser(description="Calculates Commodore checksums for a file")
+	parser.add_argument("--old",action="store_const",const="old",dest="type",default="old",help="Use Old Algorithm (default)")
+	parser.add_argument("--new",action="store_const",const="new",dest="type",help="Use New Algorithm")
+	parser.add_argument("file",type=argparse.FileType("rb"),default=None,help="Input filename")
+	parser.add_argument("checksum",nargs="?",type=str,default=None,help="Desired checksum, for inserting checksum modifier")
+	parser.add_argument("offset",nargs="?",type=str,default=None,help="Offset for checksum modifier")
+
+	parser.add_argument("-d","--dec",action="store_const",const=10,dest="base",default=16,help="Use Decimal")
+	parser.add_argument("-8","--oct",action="store_const",const=8,dest="base",help="Use Octal")
+	parser.add_argument("-2","--bin",action="store_const",const=2,dest="base",help="Use Binary")
+
+	parser.add_argument("-y","--yes",action="store_const",const=False,dest="ask",default=True,help="Don't prompt about writing to file")
+	parser.add_argument("-a","--append",action="store_const",const=True,dest="append",default=False,help="Append checksum to output (for USR files)")
+
+	parser.add_argument("-o","--output",type=str,metavar="outfile",default=None,help="Output filename (default: same as input file)")
+
+	parser.add_argument("-q","--quiet",action="store_const",const=True,dest="quiet",default=False,help="Don't output anything except the checksum")
+
+
+	opts = parser.parse_args()
+
+	if ((opts.checksum is not None and opts.offset is None) or
+		(opts.offset is not None and opts.checksum is None)
+	):
+			print("Error: Both or neither of offset and checksum are required")
+			parser.print_usage()
+			sys.exit()
+
+
+	#print(opts)
+	old = (opts.type == "old")
+
+	data = bytearray(opts.file.read())
+
+	checksum = calc_checksum(data, old)
+	opts.file.close()
+
+	disp_checksum=to_base(checksum,opts.base)
+
+	msg = f"{opts.file.name}: {disp_checksum}"
+	if opts.quiet:
+		msg = disp_checksum
 
-if sys.argv[1] == "--old":
-	old = True
-elif sys.argv[1] == "--new":
-	old = False
-else:
-	usage()
+	print(msg)
+
+	outfile = opts.output
+	if outfile is None:
+		outfile = opts.file.name
 
-filename = sys.argv[2]
+	if opts.checksum is not None and opts.offset is not None:
+
+		desired_checksum = int(opts.checksum, opts.base)
+		offset = int(opts.offset, opts.base)
 
-data = bytearray(open(filename, 'rb').read())
+		if checksum < desired_checksum:
+			data[offset] = desired_checksum - checksum
+		else:
+			data[offset] = 0xff - (checksum - desired_checksum)
 
-checksum = calc_checksum(data, old)
 
-if len(sys.argv) != 5:
-	print(filename + ": " + hex(checksum))
-	exit(0)
+		if opts.ask and not confirm(f"Write byte {offset} of '{outfile}'? "): sys.exit()
+
+		file = open(outfile, "wb")
+		file.write(data)
 
-desired_checksum = int(sys.argv[3], 16)
-offset = int(sys.argv[4], 16)
+	if opts.append:
+		# Append the checksum to the file, for e.g "&file" USR disk utilities	
+		if opts.ask and not confirm(f"Append byte to '{outfile}'? "): sys.exit()
 
-if checksum < desired_checksum:
-	data[offset] = desired_checksum - checksum
-else:
-	data[offset] = 0xff - (checksum - desired_checksum)
+		data.append(checksum)
+		file = open(outfile, "wb")
+		file.write(data)
 
-file = open(filename, "wb")
-file.write(data)
+