taikodb_hash.py

import os
import sys
import hashlib
import base64
import sqlite3

def md5(md5hash, filename):
	with open(filename, "rb") as file:
		for chunk in iter(lambda: file.read(64 * 1024), b""):
			md5hash.update(chunk)

def get_hashes(root):
	hashes = {}
	diffs = ["easy", "normal", "hard", "oni", "ura"]
	dirs = os.listdir(root)
	for dir in dirs:
		if dir.isdigit() and os.path.isdir(dir):
			files = os.listdir(os.path.join(root, dir))
			md5hash = hashlib.md5()
			if "main.tja" in files:
				md5(md5hash, os.path.join(dir, "main.tja"))
			else:
				for diff in diffs:
					if diff + ".osu" in files:
						md5(md5hash, os.path.join(root, dir, diff + ".osu"))
			hashes[dir] = base64.b64encode(md5hash.digest())[:-2]
	return hashes

def write_db(database):
	db = sqlite3.connect(database)
	hashes = get_hashes(".")
	for id in hashes:
		cur = db.cursor()
		cur.execute("update songs set hash = ? where id = ?", (hashes[id].decode(), int(id)))
		cur.close()
	db.commit()
	db.close()

if len(sys.argv) >= 2:
	write_db(sys.argv[1])
else:
	print("Please provide the path to the Taiko Web database file as the first argument")