divatool/xdudata/asset_extract.py

import functools
import multiprocessing
import os
import shutil
import subprocess
from contextlib import contextmanager
from distutils.dir_util import copy_tree
from glob import glob

from adx import extract_loop_data_from_dir

OPERATING_SYSTEM = os.name


def process_se(cache_dir, extract_dir, update_dir):
    copy_tree(
        os.path.join(cache_dir, "Android/Asset/Sound/Se"),
        os.path.join(extract_dir, "Se"),
    )

    from_dir = "Se"
    to_dir = "Se_Extracted"
    acb_prefix = "_vgmt_acb_ext"

    pool = get_pool()

    with change_cwd(extract_dir):
        os.makedirs(to_dir, exist_ok=True)

        to_del = glob("{}/**/*.wav".format(from_dir), recursive=True)
        to_del.extend(glob("{}/**/*.hca".format(from_dir), recursive=True))
        for file in to_del:
            os.remove(file)

        acb_files = [file for file in glob("{}/*.acb".format(from_dir))]
        acb_cmd = get_dotnet_cmd("acb_extract.exe")
        run_program(acb_files, cmd=acb_cmd)

        hca_files = [
            file for file in glob("{}/**/*.hca".format(from_dir), recursive=True)
        ]
        run_program(hca_files, cmd="clHCA")

        basenames = [os.path.splitext(os.path.basename(file))[0] for file in acb_files]
        char_secs = [basename.split("_", maxsplit=1) for basename in basenames]
        for cs in char_secs:
            char = cs[0]
            sec = str(cs[1]) if len(cs) == 2 else ""
            os.makedirs(os.path.join(to_dir, char, sec), exist_ok=True)
            wav_files = glob(
                "{}/{}_{}/**/*.wav".format(from_dir, acb_prefix, "_".join(cs)),
                recursive=True,
            )
            partial_ffmpeg_cmd = functools.partial(run_program, cmd="ffmpeg")
            wav_cmds = [
                [
                    "-i",
                    os.path.abspath(file),
                    "-acodec",
                    "libopus",
                    "-ab",
                    "192k",
                    "-af",
                    "aresample=48000",
                    os.path.abspath(
                        "{}/{}/{}/{}.opus".format(
                            to_dir,
                            char,
                            sec,
                            os.path.splitext(os.path.basename(file))[0],
                        )
                    ),
                    "-n",
                ]
                for file in wav_files
            ]
            pool.map(partial_ffmpeg_cmd, wav_cmds)

    copy_tree(os.path.join(extract_dir, to_dir), os.path.join(update_dir, "Se"))


def process_bgm(cache_dir, extract_dir, update_dir):
    copy_tree(
        os.path.join(cache_dir, "Android/Asset/Sound/Bgm"),
        os.path.join(extract_dir, "Bgm"),
    )

    from_dir = "Bgm"
    to_dir = "Bgm_Extracted"
    loop_file = "loop.json"

    pool = get_pool()

    with change_cwd(extract_dir):
        os.makedirs(to_dir, exist_ok=True)

        to_del = glob("{}/**/*.adx".format(from_dir), recursive=True)
        for file in to_del:
            os.remove(file)

        acb_files = [file for file in glob("{}/*.acb".format(from_dir))]
        acb_cmd = get_dotnet_cmd("acb_extract.exe")
        run_program(acb_files, cmd=acb_cmd)

        adx_files = [
            file for file in glob("{}/**/*.adx".format(from_dir), recursive=True)
        ]
        partial_ffmpeg_cmd = functools.partial(run_program, cmd="ffmpeg")
        wav_cmds = [
            [
                "-i",
                os.path.abspath(file),
                "-acodec",
                "libopus",
                "-ab",
                "192k",
                "-af",
                "aresample=48000",
                os.path.abspath(
                    "{}/{}.opus".format(
                        to_dir, os.path.splitext(os.path.basename(file))[0],
                    )
                ),
                "-n",
            ]
            for file in adx_files
        ]
        pool.map(partial_ffmpeg_cmd, wav_cmds)

        extract_loop_data_from_dir(from_dir, loop_file)
        shutil.copy2(loop_file, os.path.join(to_dir, "BgmLoop.json"))

    copy_tree(os.path.join(extract_dir, to_dir), os.path.join(update_dir, "Bgm"))


def process_voice(cache_dir, extract_dir, update_dir):
    copy_tree(
        os.path.join(cache_dir, "Android/Asset/Sound/Voice"),
        os.path.join(extract_dir, "Voice"),
    )
    subprocess.run(
        os.path.join("TODO", "update_voice.zsh"), shell=True, cwd=extract_dir
    )
    copy_tree(
        os.path.join(extract_dir, "voice_extracted"), os.path.join(update_dir, "Voice")
    )


def process_movie(cache_dir, extract_dir, update_dir):
    copy_tree(
        os.path.join(cache_dir, "Common/Asset/Movie"),
        os.path.join(extract_dir, "Movie"),
    )
    subprocess.run(
        os.path.join("TODO", "update_movie.zsh"), shell=True, cwd=extract_dir
    )
    copy_tree(
        os.path.join(extract_dir, "movie_extracted"),
        os.path.join(update_dir, "Asset/Movie"),
    )


@contextmanager
def change_cwd(new_path):
    old_cwd = os.getcwd()
    os.chdir(new_path)
    try:
        yield
    finally:
        os.chdir(old_cwd)


def get_dotnet_cmd(name):
    ext = os.path.splitext(name)[1]
    if ext != ".exe":
        raise RuntimeError("This doesn't seem like a .NET program.")
    path = shutil.which(name)
    if OPERATING_SYSTEM == "nt":
        return [path]
    return ["mono", path]


def run_program(args, cmd=None):
    if isinstance(cmd, str):
        cmd = [cmd]
    full_cmd = cmd + args
    subprocess.run(full_cmd)


def get_pool(size=multiprocessing.cpu_count()):
    pool = multiprocessing.Pool(size)
    return pool