converting a maimai FiNALE dump to simai charts

convert.py

# SPDX-License-Identifier: 0BSD
# (c) beerpsi 2022
#
# Jank ass script to convert a maimai "classic" dump into simai charts.
#
# Requires imagemagick, ffmpeg and unxwb.
#
# Put this in the data folder, install maiconverter and pathvalidate,
# provide the encryption key (MiLK and FiNALE), then run.

import argparse
import json
import os
import re
import shutil
import subprocess
from dataclasses import dataclass
from math import floor
from random import randrange
from typing import Union
from urllib.parse import parse_qs, quote, urlencode

import ffmpeg
from maiconverter.converter import sdt_to_simai
from maiconverter.maicrypt import finale_file_decrypt
from maiconverter.maisxt import MaiSxt
from pathvalidate import sanitize_filename
import requests
from tqdm import tqdm


@dataclass
class MMMusic:
    id: int
    name: str
    ver: int
    subcategory: int
    bpm: float
    sortid: int
    dress: int
    darkness: int
    mile: int
    vl: int
    event: int
    rec: int
    pv_start: float
    pv_end: float
    song_length: int
    song_ranking: int
    ad_def: int
    remaster: int
    special_pv: int
    challenge_track: int
    bonus: int
    genre_id: int
    title: str
    artist: str
    sort_jp_index: int
    sort_ex_index: int
    filename: str


@dataclass
class MMScore:
    id: int
    name: str
    lv: float
    designer_id: int
    calculated_target: int
    safe_name: str


EXTRACT_FROM_STRING_REGEX = re.compile(r"L?\"(.+)\"")
TRIM_END_BRACKET_REGEX = re.compile(r" \)( ///<.+)?$")
CHART_FILE_EXTENSIONS = ["srb", "srt", "szb", "szt", "scb", "sct", "sdb", "sdt"]
VERSIONS = [
    "01. maimai",
    "02. maimai PLUS",
    "03. GreeN",
    "04. GreeN PLUS",
    "05. ORANGE",
    "06. ORANGE PLUS",
    "07. PiNK",
    "08. PiNK PLUS",
    "09. MURASAKi",
    "10. MURASAKi PLUS",
    "11. MiLK",
    "12. MiLK PLUS",
    "13. FiNALE",
]
UTAGE_DIFFICULTIES = [
    "宴",
    "狂",
    "蔵",
    "耐",
    "蛸",
    "光",
    "星",
    "傾",
    "",
    "",
    "覚",
    "協",
    "星",
    "",
    "即",
    "撫",
]


def maidata_loads(s: str) -> "dict[str, str]":
    return parse_qs(s)


def maidata_dumps(d: "dict[str, str]") -> str:
    def maidata_key_order(kv: "tuple[str, str]") -> int:
        k = kv[0]
        if "inote_" in k:
            return 993 + int(k.replace("inote_", "")) - 1
        if "first_" in k:
            return 986 + int(k.replace("first_", "")) - 1
        if "des_" in k:
            return 979 + int(k.replace("des_", "")) - 1
        if "lv_" in k:
            return 972 + int(k.replace("lv_", "")) - 1

        KEY_ORDER = {
            "title": 0,
            "artist": 1,
            "smsg": 2,
            "des": 2,
            "first": 3,
            "wholebpm": 4,
        }
        if k in KEY_ORDER:
            return KEY_ORDER[k]
        return randrange(5, 986)
    
    items = sorted(d.items(), key=maidata_key_order)
    return "&" + urlencode(items, quote_via=quote, safe="(){}[],:#|┃/`$-^<>*@?!").replace("&", "\n&")


def read_chart(file: str, key: Union[str, bytes], bpm: float) -> MaiSxt:
    sdt = MaiSxt(bpm=bpm)

    data = ""
    if file.endswith("b"):
        data = finale_file_decrypt(file, key).decode("utf-8")
    else:
        with open(file, "r") as f:
            data = f.read()

    for line in data.split("\r\n"):
        if not line:
            continue
        if file.endswith(".srb"):
            sdt.parse_srt_line(line)
        else:
            sdt.parse_line(line)
    return sdt


def get_ver_folder(ver: int) -> str:
    if ver < 18000:
        res = VERSIONS[ver // 1000 - 10]
    elif 18000 <= ver < 19900:
        res = VERSIONS[ver // 500 - 28]
    else:
        res = VERSIONS[12]
    return res


def parse_args():
    parser = argparse.ArgumentParser(
        description="Script to convert a maimai FiNALE dump into simai charts"
    )
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        "-k",
        "--key",
        type=str,
        required=False,
        help="16 byte AES key for encrypt/decrypt (Prepend hex value with 0x)",
    )
    parser.add_argument(
        "-l",
        "--language",
        choices=["ex", "jp"],
        help="The strings file to use (changes song titles and artists)",
    )
    parser.add_argument(
        "-p",
        "--pv",
        action="store_true",
        help="Also encode PVs",
    )
    parser.add_argument("-o", "--output", default="./output", help="Output folder")
    return parser.parse_args()


def process_mmscore(
    args: argparse.Namespace,
    line: str,
    mmtext: "dict[str, str]",
    mmmusic: "dict[int, MMMusic]",
    mmfeslist: "dict[int, str]",
):
    data = line.split(",")
    meta = MMScore(
        id=int(data[0]),
        name=data[1].strip(),
        lv=float(data[2]),
        designer_id=int(data[3]),
        calculated_target=int(data[4]),
        safe_name=EXTRACT_FROM_STRING_REGEX.findall(data[5])[0].strip(),
    )
    song_id = meta.id // 100
    if song_id not in mmmusic:
        return
    song_meta = mmmusic[song_id]
    song_ver = get_ver_folder(song_meta.ver)
    song_title_sanitized = sanitize_filename(song_meta.title)

    sdt = None
    for ext in CHART_FILE_EXTENSIONS:
        filename = os.path.join("score", f"{meta.safe_name}.{ext}")
        if os.path.exists(os.path.join("score", f"{meta.safe_name}.{ext}")):
            sdt = read_chart(filename, args.key, song_meta.bpm)
            break
    if sdt is None:
        return
    simai = sdt_to_simai(sdt).export(max_den=1000)

    difficulty_id = meta.id % 100
    if difficulty_id <= 6:  # normal charts
        maidata_path = os.path.join(
            args.output, song_ver, song_title_sanitized, "maidata.txt"
        )
        with open(maidata_path, "r") as f:
            maidata = maidata_loads(f.read())

        with open(maidata_path, "w") as f:
            f.write(
                maidata_dumps(
                    {
                        **maidata,
                        **{
                            f"lv_{difficulty_id}": f'{floor(meta.lv)}{"+" if meta.lv - floor(meta.lv) >= 0.7 else ""}',
                            f"inote_{difficulty_id}": simai,
                            f"des_{difficulty_id}": mmtext.get(
                                f"RST_SCORECREATOR_{str(meta.designer_id).zfill(4)}", ""
                            ),
                        },
                    }
                )
            )
    else:  # utage
        difficulty = mmfeslist[meta.id]
        if (
            song_meta.title == "Garakuta Doll Play"
            or song_meta.title == "Wonderland Wars オープニング"
        ):  # edge case im too lazy to handle properly
            difficulty += f" NO.{meta.id % 10 + 1}"
        song_folder = os.path.join(
            args.output, song_ver, f"[{difficulty}] {song_title_sanitized}"
        )
        os.makedirs(song_folder, exist_ok=True)
        maidata = {
            "title": f"[宴] {song_meta.title}",
            "artist": song_meta.artist,
            "wholebpm": song_meta.bpm,
            "lv_7": mmfeslist[meta.id],
            "inote_7": simai,
        }
        with open(os.path.join(song_folder, "maidata.txt"), "w") as f:
            f.write(maidata_dumps(maidata))
        for file in ["bg.jpg", "track.mp3", "pv.mp4"]:
            origfile = os.path.join(args.output, song_ver, song_title_sanitized, file)
            if os.path.exists(origfile):
                shutil.copyfile(origfile, os.path.join(song_folder, file))


def encode_pv(file: str, out: str):
    (
        ffmpeg.input(file, hwaccel="dxva2")
        .filter("scale", 300, -1)
        .filter("fps", 30)
        .output(out, vcodec="h264_qsv")
        .run(quiet=True)
    )


def encode_track(padded_id: str, song_folder: str):
    subprocess.call(
        f'unxwb -d "{song_folder}" -b ./sound/{padded_id}.xsb 0 ./sound/{padded_id}.xwb',
        shell=False,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    ffmpeg.input(os.path.join(song_folder, "SDBK-.wav")).output(
        os.path.join(song_folder, "track.mp3"), ar="44100"
    ).run(quiet=True)
    os.remove(os.path.join(song_folder, "SDBK-.wav"))


def read_table(file: str, key: str) -> str:
    if file.endswith(".bin"):
        return finale_file_decrypt(file, key).decode("utf-16-le")
    else:
        with open(file, "r") as f:
            return f.read()


def main():
    args = parse_args()
    tbl_ext = "bin" if args.key else "tbl"

    for ver in VERSIONS:
        os.makedirs(os.path.join(args.output, ver), exist_ok=True)

    mmtext: dict[str, str] = {}
    mmmusic: dict[int, MMMusic] = {}
    mmfeslist: dict[int, str] = {}

    textout_file = os.path.join("tables", f"mmtextout_{args.language}.{tbl_ext}")
    print(f"[*] Reading strings from {textout_file}")
    for line in tqdm(
        [
            TRIM_END_BRACKET_REGEX.sub("", x.replace("MMTEXTOUT( ", "")).strip()
            for x in read_table(textout_file, args.key).split("\r\n")
            if x.startswith("MMTEXTOUT(")
        ]
    ):
        data = line.split('" ,L"', 2)
        mmtext[data[0][2:]] = data[1][:-1]

    print(f"[*] Reading UTAGE difficulties")
    for line in tqdm(
        [
            TRIM_END_BRACKET_REGEX.sub("", x.replace("MMFESLIST( ", "")).strip()
            for x in read_table(f"./tables/mmFesList.{tbl_ext}", args.key).split("\r\n")
            if x.startswith("MMFESLIST(")
        ]
    ):
        data = line.split(", ")
        mmfeslist[int(data[4])] = UTAGE_DIFFICULTIES[int(data[5]) - 1]

    print(f"[*] Reading music metadata, encoding music and album covers")
    for line in tqdm(
        [
            TRIM_END_BRACKET_REGEX.sub("", x.replace("MMMUSIC( ", "")).strip()
            for x in read_table(f"./tables/mmMusic.{tbl_ext}", args.key).split("\r\n")
            if x.startswith("MMMUSIC(")
        ]
    ):
        data = line.split(", ")
        meta = MMMusic(
            id=int(data[0]),
            name=data[1].strip(),
            ver=int(data[2]),
            subcategory=int(data[3]),
            bpm=float(data[4]),
            sortid=int(data[5]),
            dress=int(data[6]),
            darkness=int(data[7]),
            mile=int(data[8]),
            vl=int(data[9]),
            event=int(data[10]),
            rec=int(data[11]),
            pv_start=float(data[12]),
            pv_end=float(data[13]),
            song_length=int(data[14]),
            song_ranking=int(data[15]),
            ad_def=int(data[16]),
            remaster=int(data[17]),
            special_pv=int(data[18]),
            challenge_track=int(data[19]),
            bonus=int(data[20]),
            genre_id=int(data[21]),
            title=mmtext.get(data[22], "").strip(),
            artist=mmtext.get(data[23], "").strip(),
            sort_jp_index=int(data[24]),
            sort_ex_index=int(data[25]),
            filename=EXTRACT_FROM_STRING_REGEX.findall(data[26])[0].strip(),
        )
        mmmusic[meta.id] = meta

        ver = get_ver_folder(meta.ver)
        song_folder = os.path.join(args.output, ver, sanitize_filename(meta.title))
        os.makedirs(song_folder, exist_ok=True)

        maidata = {
            "title": meta.title,
            "artist": meta.artist,
            "wholebpm": meta.bpm,
        }
        with open(os.path.join(song_folder, "maidata.txt"), "w") as f:
            f.write(maidata_dumps(maidata))

        padded_id = str(meta.id).zfill(3)

        if not os.path.exists(os.path.join(song_folder, "bg.jpg")):
            subprocess.call(
                f'magick convert ./sprite/movie_selector/{padded_id}_mms_{meta.filename}.dds -resize 600x600 "{os.path.join(song_folder, "bg.jpg")}"',
                shell=False,
                stderr=subprocess.DEVNULL,
                stdout=subprocess.DEVNULL,
            )
        if (
            args.pv
            and not os.path.exists(os.path.join(song_folder, "pv.mp4"))
            and os.path.exists(f"./movie/{padded_id}_mmv_{meta.filename}.wmv")
        ):
            encode_pv(
                f"./movie/{padded_id}_mmv_{meta.filename}.wmv",
                os.path.join(song_folder, "pv.mp4"),
            )
        if not os.path.exists(os.path.join(song_folder, "track.mp3")):
            encode_track(padded_id, song_folder)

    print("[*] Reading and converting charts")
    for line in tqdm(
        [
            TRIM_END_BRACKET_REGEX.sub("", x.replace("MMSCORE( ", "")).strip()
            for x in (read_table(f"./tables/mmScore.{tbl_ext}", args.key).split("\r\n"))
            if x.startswith("MMSCORE(")
        ]
    ):
        process_mmscore(args, line, mmtext, mmmusic, mmfeslist)

    if args.language == "jp":
        print("[*] Updating difficulties")
        resp = requests.get("https://maimai.sega.jp/data/maimai_songs.json")
        if not resp.ok:
            print("[!] Cannot fetch https://maimai.sega.jp/data/maimai_songs.json")
            return
        songs = json.loads(resp.text)
        for song in songs:
            song_folder = os.path.join(
                args.output,
                get_ver_folder(int(song["version"])),
                sanitize_filename(song["title"]),
            )
            if not os.path.exists(os.path.join(song_folder, "maidata.txt")):
                continue

            with open(os.path.join(song_folder, "maidata.txt"), "r") as f:
                maidata = maidata_loads(f.read())

            levels = ["bas", "adv", "exp", "mas", "remas"]
            for i in range(2, 7):
                if (
                    song.get(f"lev_{levels[i - 2]}") is not None
                    and maidata.get(f"lv_{i}") is not None
                ):
                    maidata[f"lv_{i}"] = song.get(f"lev_{levels[i - 2]}")

            with open(os.path.join(song_folder, "maidata.txt"), "w") as f:
                f.write(maidata_dumps(maidata))
    else:
        print("[!] Cannot update difficulties when language is not JP.")
        print(
            "    This is because SEGA decided against localizing titles in maimai DX."
        )
        print("    I might figure this out eventually but I'm lazy.")

    print("[*] Done!")


if __name__ == "__main__":
    main()