extract_android_fmod_banks.py

from __future__ import annotations

import argparse
import shutil
import sys
import tempfile
import zipfile
from pathlib import Path

from sts2_recovery.common import find_default_gdre, run


def _collect_bank_members(zip_path: Path) -> list[zipfile.ZipInfo]:
    with zipfile.ZipFile(zip_path) as zf:
        return [info for info in zf.infolist() if info.filename.lower().endswith(".bank")]


def _extract_bank_members(zip_path: Path, members: list[zipfile.ZipInfo], output_dir: Path) -> int:
    count = 0
    with zipfile.ZipFile(zip_path) as zf:
        for info in members:
            if info.is_dir():
                continue
            name = info.filename.replace("\\", "/")
            if "/banks/" in name:
                relative = name.split("/banks/", 1)[1]
            else:
                relative = Path(name).name
            dest = output_dir / relative
            dest.parent.mkdir(parents=True, exist_ok=True)
            with zf.open(info) as src, dest.open("wb") as dst:
                shutil.copyfileobj(src, dst)
            count += 1
    return count


def _extract_pck_from_zip(zip_path: Path, temp_dir: Path) -> list[Path]:
    pck_files: list[Path] = []
    with zipfile.ZipFile(zip_path) as zf:
        for info in zf.infolist():
            if info.is_dir():
                continue
            if info.filename.lower().endswith(".pck"):
                dest = temp_dir / Path(info.filename).name
                with zf.open(info) as src, dest.open("wb") as dst:
                    shutil.copyfileobj(src, dst)
                pck_files.append(dest)
    return pck_files


def _recover_pck_banks(pck_path: Path, gdre_path: Path, output_dir: Path) -> int:
    with tempfile.TemporaryDirectory(prefix="sts2_pck_extract_") as tmpdir:
        tmpdir_path = Path(tmpdir)
        run(
            [
                str(gdre_path),
                "--headless",
                f"--recover={pck_path}",
                f"--output={tmpdir_path}",
                "--quit",
            ],
            check=True,
        )
        banks_root = tmpdir_path / "banks"
        if not banks_root.exists():
            return 0
        bank_files = list(banks_root.rglob("*.bank"))
        for bank in bank_files:
            relative = bank.relative_to(banks_root)
            dest = output_dir / relative
            dest.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(bank, dest)
        return len(bank_files)


def extract_from_package(package_path: Path, output_dir: Path, gdre_path: Path) -> int:
    bank_members = _collect_bank_members(package_path)
    if bank_members:
        return _extract_bank_members(package_path, bank_members, output_dir)

    with tempfile.TemporaryDirectory(prefix="sts2_pck_zip_") as tmpdir:
        pck_files = _extract_pck_from_zip(package_path, Path(tmpdir))
        total = 0
        for pck in pck_files:
            total += _recover_pck_banks(pck, gdre_path, output_dir)
        return total


def main() -> int:
    parser = argparse.ArgumentParser(description="Extract Android FMOD .bank files from an official package.")
    parser.add_argument("--apk", type=Path, help="Path to official Android APK")
    parser.add_argument("--obb", type=Path, help="Path to official Android OBB (zip)")
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("vendor") / "android" / "fmod" / "banks",
        help="Output directory for extracted banks",
    )
    parser.add_argument("--gdre", type=Path, help="Optional path to gdre_tools.exe")
    args = parser.parse_args()

    if args.apk is None and args.obb is None:
        print("error: provide --apk and/or --obb", file=sys.stderr)
        return 2

    output_dir = args.output.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    gdre_path = args.gdre.resolve() if args.gdre else find_default_gdre(Path(__file__).parent)
    if not gdre_path.exists():
        print(f"error: gdre_tools.exe not found at {gdre_path}", file=sys.stderr)
        return 2

    total = 0
    for package in [args.apk, args.obb]:
        if package is None:
            continue
        package_path = package.resolve()
        if not package_path.exists():
            print(f"error: package not found: {package_path}", file=sys.stderr)
            return 2
        extracted = extract_from_package(package_path, output_dir, gdre_path)
        print(f"extracted {extracted} banks from {package_path}")
        total += extracted

    if total == 0:
        print("warning: no .bank files were found in the provided package(s)")
        return 1

    print(f"done: extracted {total} banks to {output_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())