japanese-drum-game/tools/taikodb_hash.py

50 lines
1.3 KiB
Python
Raw Normal View History

2020-03-07 10:39:48 +08:00
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:
dir_path = os.path.join(root, dir)
if dir.isdigit() and os.path.isdir(dir_path):
files = os.listdir(dir_path)
md5hash = hashlib.md5()
if "main.tja" in files:
md5(md5hash, os.path.join(dir_path, "main.tja"))
else:
for diff in diffs:
if diff + ".osu" in files:
md5(md5hash, os.path.join(dir_path, diff + ".osu"))
hashes[dir] = base64.b64encode(md5hash.digest())[:-2]
return hashes
def write_db(database, songs):
db = sqlite3.connect(database)
hashes = get_hashes(songs)
added = 0
for id in hashes:
added += 1
cur = db.cursor()
cur.execute("update songs set hash = ? where id = ?", (hashes[id].decode(), int(id)))
cur.close()
db.commit()
db.close()
if added:
print("{0} hashes have been added to the database.".format(added))
else:
print("Error: No songs were found in the given directory.")
if len(sys.argv) >= 3:
write_db(sys.argv[1], sys.argv[2])
else:
print("Usage: taikodb_hash.py ../taiko.db ../public/songs")