Beat Saber の未圧縮曲ディレクトリを BMBF 用の zip ファイルに変換する Python スクリプト

beatsaber_compress_songs_for_bmbf.py

"""Beat Saber の未圧縮ディレクトリを zip 形式に変換する

Beat Saber の曲ディレクトリが以下のようになっているとき

* 親ディレクトリ
  * 曲ディレクトリ1
  * 曲ディレクトリ2
  * 曲ディレクトリ3

以下を出力する

* 曲1 ~ 3 を zip 化したもの
* 曲1 ~ 3 を Playlist 化したもの

これら出力は BMBF の Upload 画面で一括 Upload できる
"""

import argparse
import hashlib
import json
import os
import shutil
from glob import glob
from typing import List


class Song:
    """曲ディレクトリを扱うクラス
    """

    def __init__(self, song_dir: str):
        self._song_dir = song_dir
        self._info_path = os.path.join(song_dir, "info.dat")
        with open(self._info_path) as f:
            self._info_dat = json.load(f)

    def compress(self, out_dir: str):
        """曲を zip 化する
        """
        # BMBF が取り扱う曲の zip ファイルは、ディレクトリ自体を含まず
        # 曲データが直接トップレベルに格納されるよう圧縮する必要がある。
        # ここでは標準でそのような挙動となる shutil.make_archive を使っている
        shutil.make_archive(os.path.join(out_dir, self.hash),
                            "zip", self._song_dir)

    @property
    def hash(self):
        """曲のハッシュ値

        Beat Saber Mod (BMBF) で取り扱う譜面は、 (曲データのハッシュ値).zip というファイル名になっていることを想定している。
        そのため、未圧縮の曲を zip 化するためにハッシュ値を計算する必要がある。

        曲ハッシュ値は SHA1 で計算されており、info.dat と各難易度の譜面データを cat したものに対して
        SHA1 を計算することで得られる
        """
        if not hasattr(self, '_song_hash'):
            m = hashlib.sha1()
            with open(self._info_path, "rb") as f:
                m.update(f.read())

            for difficulty_set in self._info_dat["_difficultyBeatmapSets"]:
                for difficulty in difficulty_set["_difficultyBeatmaps"]:
                    with open(os.path.join(self._song_dir, difficulty["_beatmapFilename"]), "rb") as f:
                        m.update(f.read())

            self._song_hash = m.hexdigest()
        return self._song_hash

    @property
    def song_name(self):
        return self._info_dat["_songName"]

    @property
    def song_author(self):
        return self._info_dat["_songAuthorName"]

    @property
    def level_author(self):
        return self._info_dat["_levelAuthorName"]


def get_song_dirs(parent_dir: str) -> List[str]:
    """曲のディレクトリ一覧を取得する

    曲は以下のようにディレクトリに保存されていると仮定している

    parent_dir
    - 曲1
    - 曲2
    - 曲3

    Args:
        parent_dir (str): 曲を含むディレクトリ

    Returns:
        list[str]: 曲ディレクトリのパスの配列
    """
    return glob(os.path.join(parent_dir, "*"))


def gen_playlist(title: str, author: str, description: str, songs: List[Song]) -> dict:
    """Beat Saber の PlaylistManager Mod で読み込めるプレイリストを作成する

    Args:
        title (str):
        author (str):
        description (str):
        songs: (list[Song]): 

    Returns:
        dict[str, str]: PlaylistManager Mod 用のプレイリスト
    """
    return {
        "playlistTitle": title,
        "playlistAuthor": author,
        "playlistDescription": description,
        "songs": [{"hash": song.hash, "name": song.song_name} for song in songs]
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--target-dir', required=True, type=str)
    parser.add_argument('--playlist-name', required=True, type=str)
    parser.add_argument('--playlist-author', required=True, type=str)
    parser.add_argument('--playlist-description', required=True, type=str)
    args = parser.parse_args()

    songs = [Song(song_dir) for song_dir in get_song_dirs(args.target_dir)]
    songs = sorted(songs, key=lambda s: s.song_name)
    playlist = gen_playlist(args.playlist_name, args.playlist_author,
                            args.playlist_description, songs)
    with open("./out/playlist.json", "w") as f:
        json.dump(playlist, f, indent=2)
    for song in songs:
        song.compress("./out")


if __name__ == '__main__':
    main()