Scratchpad cache rebuild script for EXEPoN/EXELoN

rebuild_cache.py

import argparse
import dataclasses
import pathlib
import struct
import zipfile


@dataclasses.dataclass
class QuestData:
    scenario_num: int
    prereq: int
    quest_num: int


SP_HDR_SIZE = 0x40

PON_QUESTS = [
    QuestData(scenario_num=1, prereq=0, quest_num=1),
    QuestData(scenario_num=1, prereq=0, quest_num=2),
    QuestData(scenario_num=1, prereq=0, quest_num=3),
    QuestData(scenario_num=1, prereq=0, quest_num=4),
    QuestData(scenario_num=1, prereq=0, quest_num=5),
    QuestData(scenario_num=1, prereq=0, quest_num=6),
    QuestData(scenario_num=1, prereq=0, quest_num=7),
    QuestData(scenario_num=2, prereq=0, quest_num=8),
    QuestData(scenario_num=2, prereq=0, quest_num=9),
    QuestData(scenario_num=2, prereq=0, quest_num=10),
    QuestData(scenario_num=2, prereq=0, quest_num=11),
    QuestData(scenario_num=2, prereq=0, quest_num=12),
    QuestData(scenario_num=2, prereq=0, quest_num=13),
    QuestData(scenario_num=2, prereq=0, quest_num=14),
    QuestData(scenario_num=3, prereq=0, quest_num=15),
    QuestData(scenario_num=3, prereq=0, quest_num=16),
    QuestData(scenario_num=3, prereq=0, quest_num=17),
    QuestData(scenario_num=3, prereq=0, quest_num=18),
    QuestData(scenario_num=3, prereq=12, quest_num=19),
    QuestData(scenario_num=3, prereq=7, quest_num=20),
    QuestData(scenario_num=3, prereq=14, quest_num=21),
    QuestData(scenario_num=4, prereq=0, quest_num=22),
    QuestData(scenario_num=4, prereq=0, quest_num=23),
    QuestData(scenario_num=4, prereq=0, quest_num=24),
    QuestData(scenario_num=4, prereq=0, quest_num=25),
    QuestData(scenario_num=4, prereq=19, quest_num=26),
    QuestData(scenario_num=4, prereq=13, quest_num=27),
    QuestData(scenario_num=4, prereq=21, quest_num=28),
    QuestData(scenario_num=5, prereq=0, quest_num=29),
    QuestData(scenario_num=5, prereq=0, quest_num=30),
    QuestData(scenario_num=5, prereq=0, quest_num=31),
    QuestData(scenario_num=5, prereq=18, quest_num=32),
    QuestData(scenario_num=5, prereq=26, quest_num=33),
    QuestData(scenario_num=5, prereq=20, quest_num=34),
    QuestData(scenario_num=5, prereq=28, quest_num=35),
    QuestData(scenario_num=6, prereq=0, quest_num=36),
    QuestData(scenario_num=6, prereq=0, quest_num=37),
    QuestData(scenario_num=6, prereq=0, quest_num=38),
    QuestData(scenario_num=6, prereq=0, quest_num=39),
    QuestData(scenario_num=6, prereq=17, quest_num=40),
    QuestData(scenario_num=6, prereq=31, quest_num=41),
    QuestData(scenario_num=6, prereq=35, quest_num=42),
    QuestData(scenario_num=7, prereq=27, quest_num=43),
    QuestData(scenario_num=7, prereq=0, quest_num=44),
    QuestData(scenario_num=7, prereq=0, quest_num=45),
    QuestData(scenario_num=7, prereq=32, quest_num=46),
    QuestData(scenario_num=7, prereq=34, quest_num=47),
    QuestData(scenario_num=7, prereq=0, quest_num=48),
    QuestData(scenario_num=7, prereq=42, quest_num=49),
    QuestData(scenario_num=8, prereq=9, quest_num=50),
    QuestData(scenario_num=8, prereq=0, quest_num=51),
    QuestData(scenario_num=8, prereq=0, quest_num=52),
    QuestData(scenario_num=8, prereq=38, quest_num=53),
    QuestData(scenario_num=8, prereq=0, quest_num=54),
    QuestData(scenario_num=8, prereq=33, quest_num=55),
    QuestData(scenario_num=8, prereq=49, quest_num=56),
    QuestData(scenario_num=9, prereq=0, quest_num=57),
    QuestData(scenario_num=9, prereq=0, quest_num=58),
    QuestData(scenario_num=9, prereq=46, quest_num=59),
    QuestData(scenario_num=9, prereq=43, quest_num=60),
    QuestData(scenario_num=9, prereq=47, quest_num=61),
    QuestData(scenario_num=9, prereq=56, quest_num=62),
    QuestData(scenario_num=9, prereq=62, quest_num=63),
]
LON_QUESTS = [
    QuestData(scenario_num=2, prereq=0, quest_num=1),
    QuestData(scenario_num=2, prereq=0, quest_num=2),
    QuestData(scenario_num=2, prereq=0, quest_num=3),
    QuestData(scenario_num=2, prereq=0, quest_num=4),
    QuestData(scenario_num=2, prereq=0, quest_num=5),
    QuestData(scenario_num=2, prereq=0, quest_num=6),
    QuestData(scenario_num=2, prereq=0, quest_num=7),
    QuestData(scenario_num=2, prereq=0, quest_num=8),
    QuestData(scenario_num=3, prereq=0, quest_num=9),
    QuestData(scenario_num=3, prereq=0, quest_num=10),
    QuestData(scenario_num=3, prereq=0, quest_num=11),
    QuestData(scenario_num=3, prereq=0, quest_num=12),
    QuestData(scenario_num=3, prereq=4, quest_num=13),
    QuestData(scenario_num=3, prereq=6, quest_num=14),
    QuestData(scenario_num=3, prereq=7, quest_num=15),
    QuestData(scenario_num=3, prereq=8, quest_num=16),
    QuestData(scenario_num=4, prereq=0, quest_num=17),
    QuestData(scenario_num=4, prereq=0, quest_num=18),
    QuestData(scenario_num=4, prereq=0, quest_num=19),
    QuestData(scenario_num=4, prereq=0, quest_num=20),
    QuestData(scenario_num=4, prereq=5, quest_num=21),
    QuestData(scenario_num=4, prereq=14, quest_num=22),
    QuestData(scenario_num=4, prereq=15, quest_num=23),
    QuestData(scenario_num=4, prereq=16, quest_num=24),
    QuestData(scenario_num=5, prereq=0, quest_num=25),
    QuestData(scenario_num=5, prereq=0, quest_num=26),
    QuestData(scenario_num=5, prereq=0, quest_num=27),
    QuestData(scenario_num=5, prereq=13, quest_num=28),
    QuestData(scenario_num=5, prereq=21, quest_num=29),
    QuestData(scenario_num=5, prereq=22, quest_num=30),
    QuestData(scenario_num=5, prereq=23, quest_num=31),
    QuestData(scenario_num=5, prereq=24, quest_num=32),
    QuestData(scenario_num=6, prereq=0, quest_num=33),
    QuestData(scenario_num=6, prereq=0, quest_num=34),
    QuestData(scenario_num=6, prereq=0, quest_num=35),
    QuestData(scenario_num=6, prereq=0, quest_num=36),
    QuestData(scenario_num=6, prereq=12, quest_num=37),
    QuestData(scenario_num=6, prereq=30, quest_num=38),
    QuestData(scenario_num=6, prereq=31, quest_num=39),
    QuestData(scenario_num=6, prereq=32, quest_num=40),
    QuestData(scenario_num=7, prereq=0, quest_num=41),
    QuestData(scenario_num=7, prereq=0, quest_num=42),
    QuestData(scenario_num=7, prereq=0, quest_num=43),
    QuestData(scenario_num=7, prereq=28, quest_num=44),
    QuestData(scenario_num=7, prereq=29, quest_num=45),
    QuestData(scenario_num=7, prereq=38, quest_num=46),
    QuestData(scenario_num=7, prereq=39, quest_num=47),
    QuestData(scenario_num=7, prereq=40, quest_num=48),
    QuestData(scenario_num=8, prereq=0, quest_num=49),
    QuestData(scenario_num=8, prereq=0, quest_num=50),
    QuestData(scenario_num=8, prereq=0, quest_num=51),
    QuestData(scenario_num=8, prereq=37, quest_num=52),
    QuestData(scenario_num=8, prereq=45, quest_num=53),
    QuestData(scenario_num=8, prereq=46, quest_num=54),
    QuestData(scenario_num=8, prereq=47, quest_num=55),
    QuestData(scenario_num=8, prereq=48, quest_num=56),
    QuestData(scenario_num=9, prereq=0, quest_num=57),
    QuestData(scenario_num=9, prereq=0, quest_num=58),
    QuestData(scenario_num=9, prereq=0, quest_num=59),
    QuestData(scenario_num=9, prereq=0, quest_num=60),
    QuestData(scenario_num=9, prereq=44, quest_num=61),
    QuestData(scenario_num=9, prereq=54, quest_num=62),
    QuestData(scenario_num=9, prereq=55, quest_num=63),
    QuestData(scenario_num=9, prereq=56, quest_num=64),
]


arg_parser = argparse.ArgumentParser()
arg_parser.add_argument("sp_file", type=pathlib.Path)
arg_parser.add_argument("jar_file", type=pathlib.Path)
arg_parser.add_argument("game", type=str.lower, choices=["pon", "lon"])
args = arg_parser.parse_args()
sp_file: pathlib.Path = args.sp_file
jar_file: pathlib.Path = args.jar_file
pon: bool = args.game == "pon"
lon: bool = not pon

quests = PON_QUESTS if pon else LON_QUESTS
dat_file_count = 37 if pon else 42
save_offset = 300 if pon else 400
sp_file_offset = 0xD44 if pon else 0xDA8
sys_dat_offset = 0x64000 if pon else 0x64000
m_offset = 0 if pon else 0
quest_e_offset = 4424 if pon else 5120
quest_s_offset = 5448 if pon else 6144
scenario_e_offset = 8520 if pon else 10240
scenario_s_offset = 13640 if pon else 15360
quest_all_offset = 34120 if pon else 39360
quest_all_count = 10 if pon else 7
quest_m_offset = None if pon else 41460
dat_net_data_idx = 0 if pon else 0
dat_quest_no_idx = 38 if pon else 64
dat_quest_e_size_idx = 40 if pon else 66
dat_quest_s_size_idx = 41 if pon else 67
dat_scenario_e_size_idx = 42 if pon else 68
dat_scenario_s_0_size_idx = 43 if pon else 69
dat_cache_addr_idx = 48 if pon else 48
dat_scenario_m_0_size_idx = None if pon else 75
dat_scenario_m_count = None if pon else 5
dat_scenario_s_count = 5 if pon else 6
save_scenario_num_idx = 355 if pon else 444
save_quest_num_idx = 360 if pon else 449
save_qflag_idx = 423 if pon else 515

with open(sp_file, "rb+") as sp, zipfile.ZipFile(jar_file) as jar_zip:
    # Read network data sizes
    dat_files = []
    sp.seek(SP_HDR_SIZE + sp_file_offset)
    for _ in range(dat_file_count):
        dat_files.append(struct.unpack(">I", sp.read(4))[0])

    # Compute file addresses
    file_addrs = []
    file_addrs.append(sp_file_offset + dat_file_count * 4)
    for i in range(1, dat_file_count):
        file_addrs.append(file_addrs[-1] + dat_files[i - 1])

    # Set cache start address
    cache_addr = file_addrs[-1] + dat_files[-1]

    # Erase cache variables
    sp.seek(SP_HDR_SIZE + dat_net_data_idx * 0x4)
    sp.write(bytes(4 * dat_file_count))
    sp.seek(SP_HDR_SIZE + dat_quest_e_size_idx * 0x4)
    sp.write(bytes(4))
    sp.seek(SP_HDR_SIZE + dat_quest_s_size_idx * 0x4)
    sp.write(bytes(4))
    sp.seek(SP_HDR_SIZE + dat_scenario_e_size_idx * 0x4)
    sp.write(bytes(4))
    sp.seek(SP_HDR_SIZE + dat_scenario_s_0_size_idx * 0x4)
    sp.write(bytes(4 * dat_scenario_s_count))
    sp.seek(SP_HDR_SIZE + dat_cache_addr_idx * 0x4)
    sp.write(bytes(4))
    if lon:
        sp.seek(SP_HDR_SIZE + dat_scenario_m_0_size_idx * 0x4)
        sp.write(bytes(4 * dat_scenario_m_count))
    quest_e_size = 0
    quest_s_size = 0
    scenario_e_size = 0
    scenario_s_sizes = [0] * dat_scenario_s_count
    scenario_m_sizes = [0] * dat_scenario_m_count if lon else []

    # Erase cache
    sp.seek(SP_HDR_SIZE + cache_addr)
    while sp.tell() < SP_HDR_SIZE + sys_dat_offset:
        sp.write(b"\0")

    # Load current scenario and quest from save
    sp.seek(SP_HDR_SIZE + save_offset + save_scenario_num_idx * 4)
    (scenario_num,) = struct.unpack(">I", sp.read(4))
    sp.seek(SP_HDR_SIZE + save_offset + save_quest_num_idx * 4)
    (quest_num_active,) = struct.unpack(">I", sp.read(4))
    sp.seek(SP_HDR_SIZE + dat_quest_no_idx * 4)
    (quest_num_cached,) = struct.unpack(">I", sp.read(4))
    sp.seek(SP_HDR_SIZE + save_offset + save_qflag_idx * 4)
    quest_flags = list(struct.unpack(">BBBBBBBB", sp.read(8)))
    print(
        f"Caching scenario {scenario_num}, quest {quest_num_cached} cached, quest {quest_num_active} active"
    )

    # Cache m.dat
    m_addr = cache_addr + m_offset
    # Restore scenario m.dat
    if pon:
        # Use scenario 1 as base
        m = jar_zip.read(f"data/scenario/{1}/m.dat")
        sp.seek(SP_HDR_SIZE + m_addr)
        sp.write(m)
        # Overlap current scenario
        if scenario_num != 1:
            m = jar_zip.read(f"data/scenario/{scenario_num}/m.dat")
            sp.seek(SP_HDR_SIZE + m_addr)
            sp.write(m[:836])
            for i in range(69):
                sp.seek(SP_HDR_SIZE + m_addr + 836 + i * 48 + 32)
                sp.write(m[836 + i * 12 : 836 + (i + 1) * 12])
        for i in range(20):
            # This data is written when you finish or abort a quest
            # If you haven't finished or aborted a quest yet, it may be all 0 instead
            sp.seek(SP_HDR_SIZE + m_addr + i * 36 + 12)
            sp.write(struct.pack(">BBBB", 0, 0, 0, 0))
            sp.seek(SP_HDR_SIZE + m_addr + i * 36 + 28)
            sp.write(struct.pack(">BBBB", 31, 31, 31, 31))
    if lon:
        # While not strictly necessary, this mimics how consecutive scenarios are cached
        for i in range(1, scenario_num + 1):
            m = jar_zip.read(f"data/scenario/{i}/m.dat")
            sp.seek(SP_HDR_SIZE + m_addr)
            sp.write(m)
            for j in range(dat_scenario_m_count):
                (scenario_m_sizes[j],) = struct.unpack_from(">I", m, j * 4)
    if pon and quest_num_active != 0:
        # Restore quest m.dat
        m = jar_zip.read(f"data/quest/{quest_num_active}/m.dat")
        for i in range(8):
            x = m[i * 8 + 3]
            if x < 100:
                sp.seek(SP_HDR_SIZE + m_addr + 836 + x * 48 + 44)
                sp.write(m[i * 8 + 4 : i * 8 + 4 + 4])
    if lon and quest_num_cached != 0:
        m = jar_zip.read(f"data/quest/{quest_num_cached}/m.dat")
        quest_m_addr = cache_addr + quest_m_offset
        sp.seek(SP_HDR_SIZE + quest_m_addr)
        sp.write(m)

    if quest_num_cached != 0:
        # Restore quest e.jar
        quest_e_addr = cache_addr + quest_e_offset
        e = jar_zip.read(f"data/quest/{quest_num_cached}/e.jar")
        sp.seek(SP_HDR_SIZE + quest_e_addr)
        sp.write(e)
        quest_e_size = len(e)

        # Restore quest s.jar
        quest_s_addr = cache_addr + quest_s_offset
        s = jar_zip.read(f"data/quest/{quest_num_cached}/s.jar")
        sp.seek(SP_HDR_SIZE + quest_s_addr)
        sp.write(s)
        quest_s_size = len(s)

    # While not strictly necessary, this mimics how consecutive scenarios are cached
    for i in range(1, scenario_num + 1):
        scenario_e_addr = cache_addr + scenario_e_offset
        e = jar_zip.read(f"data/scenario/{i}/e.jar")
        sp.seek(SP_HDR_SIZE + scenario_e_addr)
        sp.write(e)
    scenario_e_size = len(e)

    # Restore scenario s.dat
    # While not strictly necessary, this mimics how consecutive scenarios are cached
    for i in range(1, scenario_num + 1):
        scenario_s_addr = cache_addr + scenario_s_offset
        s = jar_zip.read(f"data/scenario/{i}/s.dat")
        sp.seek(SP_HDR_SIZE + scenario_s_addr)
        sp.write(s[dat_scenario_s_count * 4 :])
    for i in range(dat_scenario_s_count):
        (scenario_s_sizes[i],) = struct.unpack(">I", s[i * 4 : (i + 1) * 4])

    # Generate quests list
    available_quests = []
    for quest in quests:
        if (
            scenario_num >= quest.scenario_num
            and quest_flags[(quest.quest_num - 1) // 8]
            >> (7 - (quest.quest_num - 1) % 8)
            == 0
        ) and (
            quest.prereq == 0
            or quest_flags[(quest.prereq - 1) // 8] >> (7 - (quest.prereq - 1) % 8) != 0
        ):
            available_quests.append(quest.quest_num)
    available_quests = available_quests[:quest_all_count]
    quest_text = jar_zip.read(f"data/quest/quest.dat")
    quest_all_addr = cache_addr + quest_all_offset
    sp.seek(SP_HDR_SIZE + quest_all_addr)
    for quest in available_quests:
        sp.write(quest_text[(quest - 1) * 360 + 60 : (quest - 1) * 360 + 360])

    # Write cache variables
    sp.seek(SP_HDR_SIZE + dat_net_data_idx * 0x4)
    for dat_file in dat_files:
        sp.write(struct.pack(">I", dat_file))
    sp.seek(SP_HDR_SIZE + dat_quest_e_size_idx * 0x4)
    sp.write(struct.pack(">I", quest_e_size))
    sp.seek(SP_HDR_SIZE + dat_quest_s_size_idx * 0x4)
    sp.write(struct.pack(">I", quest_s_size))
    sp.seek(SP_HDR_SIZE + dat_scenario_e_size_idx * 0x4)
    sp.write(struct.pack(">I", scenario_e_size))
    sp.seek(SP_HDR_SIZE + dat_scenario_s_0_size_idx * 0x4)
    for i in range(dat_scenario_s_count):
        sp.write(struct.pack(">I", scenario_s_sizes[i]))
    if lon:
        sp.seek(SP_HDR_SIZE + dat_scenario_m_0_size_idx * 0x4)
        for i in range(dat_scenario_m_count):
            sp.write(struct.pack(">I", scenario_m_sizes[i]))
    sp.seek(SP_HDR_SIZE + dat_cache_addr_idx * 0x4)
    sp.write(struct.pack(">I", cache_addr))