sm64

A Super Mario 64 decompilation
Log | Files | Refs | README | LICENSE

extract_assets.py (9572B)


      1 #!/usr/bin/env python3
      2 import sys
      3 import os
      4 import json
      5 
      6 
      7 def read_asset_map():
      8     with open("assets.json") as f:
      9         ret = json.load(f)
     10     return ret
     11 
     12 
     13 def read_local_asset_list(f):
     14     if f is None:
     15         return []
     16     ret = []
     17     for line in f:
     18         ret.append(line.strip())
     19     return ret
     20 
     21 
     22 def asset_needs_update(asset, version):
     23     if version <= 6 and asset in ["actors/king_bobomb/king_bob-omb_eyes.rgba16.png", "actors/king_bobomb/king_bob-omb_hand.rgba16.png"]:
     24         return True
     25     if version <= 5 and asset == "textures/spooky/bbh_textures.00800.rgba16.png":
     26         return True
     27     if version <= 4 and asset in ["textures/mountain/ttm_textures.01800.rgba16.png", "textures/mountain/ttm_textures.05800.rgba16.png"]:
     28         return True
     29     if version <= 3 and asset == "textures/cave/hmc_textures.01800.rgba16.png":
     30         return True
     31     if version <= 2 and asset == "textures/inside/inside_castle_textures.09000.rgba16.png":
     32         return True
     33     if version <= 1 and asset.endswith(".m64"):
     34         return True
     35     if version <= 0 and asset.endswith(".aiff"):
     36         return True
     37     return False
     38 
     39 
     40 def remove_file(fname):
     41     os.remove(fname)
     42     print("deleting", fname)
     43     try:
     44         os.removedirs(os.path.dirname(fname))
     45     except OSError:
     46         pass
     47 
     48 
     49 def clean_assets(local_asset_file):
     50     assets = set(read_asset_map().keys())
     51     assets.update(read_local_asset_list(local_asset_file))
     52     for fname in list(assets) + [".assets-local.txt"]:
     53         if fname.startswith("@"):
     54             continue
     55         try:
     56             remove_file(fname)
     57         except FileNotFoundError:
     58             pass
     59 
     60 
     61 def main():
     62     # In case we ever need to change formats of generated files, we keep a
     63     # revision ID in the local asset file.
     64     new_version = 7
     65 
     66     try:
     67         local_asset_file = open(".assets-local.txt")
     68         local_asset_file.readline()
     69         local_version = int(local_asset_file.readline().strip())
     70     except Exception:
     71         local_asset_file = None
     72         local_version = -1
     73 
     74     langs = sys.argv[1:]
     75     if langs == ["--clean"]:
     76         clean_assets(local_asset_file)
     77         sys.exit(0)
     78 
     79     all_langs = ["jp", "us", "eu", "sh", "cn"]
     80     if not langs or not all(a in all_langs for a in langs):
     81         langs_str = " ".join("[" + lang + "]" for lang in all_langs)
     82         print("Usage: " + sys.argv[0] + " " + langs_str)
     83         print("For each version, baserom.<version>.z64 must exist")
     84         sys.exit(1)
     85 
     86     asset_map = read_asset_map()
     87     all_assets = []
     88     any_missing_assets = False
     89     for asset, data in asset_map.items():
     90         if asset.startswith("@"):
     91             continue
     92         if os.path.isfile(asset):
     93             all_assets.append((asset, data, True))
     94         else:
     95             all_assets.append((asset, data, False))
     96             if not any_missing_assets and any(lang in data[-1] for lang in langs):
     97                 any_missing_assets = True
     98 
     99     if not any_missing_assets and local_version == new_version:
    100         # Nothing to do, no need to read a ROM. For efficiency we don't check
    101         # the list of old assets either.
    102         return
    103 
    104     # Late imports (to optimize startup perf)
    105     import subprocess
    106     import hashlib
    107     import tempfile
    108     from collections import defaultdict
    109 
    110     new_assets = {a[0] for a in all_assets}
    111 
    112     previous_assets = read_local_asset_list(local_asset_file)
    113     if local_version == -1:
    114         # If we have no local asset file, we assume that files are version
    115         # controlled and thus up to date.
    116         local_version = new_version
    117 
    118     # Create work list
    119     todo = defaultdict(lambda: [])
    120     for (asset, data, exists) in all_assets:
    121         # Leave existing assets alone if they have a compatible version.
    122         if exists and not asset_needs_update(asset, local_version):
    123             continue
    124 
    125         meta = data[:-2]
    126         size, positions = data[-2:]
    127         for lang, pos in positions.items():
    128             mio0 = None if len(pos) == 1 else pos[0]
    129             pos = pos[-1]
    130             if lang in langs:
    131                 todo[(lang, mio0)].append((asset, pos, size, meta))
    132                 break
    133 
    134     # Load ROMs
    135     roms = {}
    136     for lang in langs:
    137         fname = "baserom." + lang + ".z64"
    138         try:
    139             with open(fname, "rb") as f:
    140                 roms[lang] = f.read()
    141         except Exception as e:
    142             print("Failed to open " + fname + "! " + str(e))
    143             sys.exit(1)
    144         sha1 = hashlib.sha1(roms[lang]).hexdigest()
    145         with open("sm64." + lang + ".sha1", "r") as f:
    146             expected_sha1 = f.read().split()[0]
    147         if sha1 != expected_sha1:
    148             print(
    149                 fname
    150                 + " has the wrong hash! Found "
    151                 + sha1
    152                 + ", expected "
    153                 + expected_sha1
    154             )
    155             sys.exit(1)
    156 
    157     # Make sure tools exist
    158     subprocess.check_call(
    159         ["make", "-s", "-C", "tools/sm64tools/", "n64graphics", "mio0"]
    160     )
    161 
    162     subprocess.check_call(
    163         ["make", "-s", "-C", "tools/", "skyconv", "aifc_decode"]
    164     )
    165 
    166     # Go through the assets in roughly alphabetical order (but assets in the same
    167     # mio0 file still go together).
    168     keys = sorted(list(todo.keys()), key=lambda k: todo[k][0][0])
    169 
    170     # Import new assets
    171     for key in keys:
    172         assets = todo[key]
    173         lang, mio0 = key
    174         if mio0 == "@sound":
    175             rom = roms[lang]
    176             args = [
    177                 "python3",
    178                 "tools/disassemble_sound.py",
    179                 "baserom." + lang + ".z64",
    180             ]
    181             def append_args(key):
    182                 sound_ver = "sh" if lang == "cn" else lang
    183                 size, locs = asset_map["@sound " + key + " " + sound_ver]
    184                 offset = locs[lang][0]
    185                 args.append(str(offset))
    186                 args.append(str(size))
    187             append_args("ctl")
    188             append_args("tbl")
    189             if lang in ("sh", "cn"):
    190                 args.append("--shindou-headers")
    191                 append_args("ctl header")
    192                 append_args("tbl header")
    193             args.append("--only-samples")
    194             for (asset, pos, size, meta) in assets:
    195                 print("extracting", asset)
    196                 args.append(asset + ":" + str(pos))
    197             subprocess.run(args, check=True)
    198             continue
    199 
    200         if mio0 is not None:
    201             image = subprocess.run(
    202                 [
    203                     "./tools/sm64tools/mio0",
    204                     "-d",
    205                     "-o",
    206                     str(mio0),
    207                     "baserom." + lang + ".z64",
    208                     "-",
    209                 ],
    210                 check=True,
    211                 stdout=subprocess.PIPE,
    212             ).stdout
    213         else:
    214             image = roms[lang]
    215 
    216         for (asset, pos, size, meta) in assets:
    217             print("extracting", asset)
    218             input = image[pos : pos + size]
    219             os.makedirs(os.path.dirname(asset), exist_ok=True)
    220             if asset.endswith(".png"):
    221                 png_file = tempfile.NamedTemporaryFile(prefix="asset", delete=False)
    222                 try:
    223                     png_file.write(input)
    224                     png_file.flush()
    225                     png_file.close()
    226                     if asset.startswith("textures/skyboxes/") or asset.startswith("levels/ending/cake"):
    227                         if asset.startswith("textures/skyboxes/"):
    228                             imagetype = "sky"
    229                         else:
    230                             imagetype = "cake" + ("-cn" if "cn" in asset else "-eu" if "eu" in asset else "")
    231                         print(imagetype, png_file.name, asset)
    232                         subprocess.run(
    233                             [
    234                                 "./tools/skyconv",
    235                                 "--type",
    236                                 imagetype,
    237                                 "--combine",
    238                                 png_file.name,
    239                                 asset,
    240                             ],
    241                             check=True,
    242                         )
    243                     else:
    244                         w, h = meta
    245                         fmt = asset.split(".")[-2]
    246                         subprocess.run(
    247                             [
    248                                 "./tools/sm64tools/n64graphics",
    249                                 "-e",
    250                                 png_file.name,
    251                                 "-g",
    252                                 asset,
    253                                 "-f",
    254                                 fmt,
    255                                 "-w",
    256                                 str(w),
    257                                 "-h",
    258                                 str(h),
    259                             ],
    260                             check=True,
    261                         )
    262                 finally:
    263                     png_file.close()
    264                     os.remove(png_file.name)
    265             else:
    266                 with open(asset, "wb") as f:
    267                     f.write(input)
    268 
    269     # Remove old assets
    270     for asset in previous_assets:
    271         if asset not in new_assets:
    272             try:
    273                 remove_file(asset)
    274             except FileNotFoundError:
    275                 pass
    276 
    277     # Replace the asset list
    278     output = "\n".join(
    279         [
    280             "# This file tracks the assets currently extracted by extract_assets.py.",
    281             str(new_version),
    282             *sorted(list(new_assets)),
    283             "",
    284         ]
    285     )
    286     with open(".assets-local.txt", "w") as f:
    287         f.write(output)
    288 
    289 
    290 main()