dspatch.py/NDSCard/shutil.py
2025-04-27 19:12:29 -07:00

262 lines
9.5 KiB
Python

import os
import os.path
from .NDSCard import NDSCard, NDSCardHeader, NitroFooter, NDSCardFNT
from .NDSCard import NDSCardBanner, NDSCardFAT
from .NitroFS import FileAllocationEntry, DirectoryTableEntry, EntryNameTableDirectoryEntry
from .NitroFS import EntryNameTableFileEntry, EntryNameTableEndOfDirectoryEntry
from .Utils import align
class Unpacker:
"""Callable for unpacking Nintendo DS ROMs. This is meant to be used with
`shutil.unpack_archive`. This does not currently support unpacking ROMs that use overlays.
"""
def __call__(self, *args, **kwargs):
card = NDSCard(args[0])
out_dir = args[1]
out_path = lambda *s: os.path.join(out_dir, *s)
if card.main_ovt or card.sub_ovt:
raise NotImplementedError("ROMs with OVTs are not currently supported")
with open(out_path('header.bin'), 'wb') as fp:
fp.write(card.header.encode())
with open(out_path('arm9.bin'), 'wb') as fp:
fp.write(card.main_rom)
with open(out_path('arm7.bin'), 'wb') as fp:
fp.write(card.sub_rom)
if hasattr(card, 'static_footer'):
with open(out_path('static_footer.bin'), 'wb') as fp:
fp.write(card.static_footer.encode())
if hasattr(card, 'banner'):
with open(out_path('banner.bin'), 'wb') as fp:
fp.write(card.banner.encode())
if hasattr(card, 'rsa_sig'):
with open(out_path('rsa_sig.bin'), 'wb') as fp:
fp.write(card.rsa_sig)
# NitroFS
dirnames = [None] * len(card.fnt.directory_table)
dirnames[0] = 'data/' # root directory name
filenames = [None] * len(card.fat)
Unpacker._get_files(
dirnames, filenames, card.fnt.directory_table,
card.fnt.entry_name_table, 0
)
for dirname_elems in (name.split('/') for name in dirnames):
os.makedirs(out_path(*dirname_elems), exist_ok=True)
for file_id in range(len(filenames)):
with open(out_path(filenames[file_id]), 'wb') as fp:
fp.write(card.file_data[file_id])
@staticmethod
def _get_files(
dirnames: list[str],
filenames: list[str],
dir_table: NDSCardFNT.DirectoryTable,
names_table: NDSCardFNT.EntryNameTable,
dir_id: int
):
"""Get all directory names and filenames info into `dirnames` and `filenames`"""
file_id = dir_table[dir_id].file_id
for entry in names_table[dir_id]:
if isinstance(entry, EntryNameTableEndOfDirectoryEntry):
break
if isinstance(entry, EntryNameTableDirectoryEntry):
subdir_id = entry.directory_id ^ 0xf000
parent_dir_id = dir_table[subdir_id].parent_id ^ 0xf000
subdir_name = dirnames[parent_dir_id] + entry.entry_name + '/'
dirnames[subdir_id] = subdir_name
Unpacker._get_files(
dirnames, filenames, dir_table, names_table, subdir_id
)
else: # EntryNameTableFileEntry
path = dirnames[dir_id] + entry.entry_name
filenames[file_id] = path
file_id += 1
class Archiver:
"""Callable for repacking Nintendo DS ROMs that were unpacked by `NDSCard.Unpacker`. This is
meant to be used with `shutil.make_archive`. This does not currently suppport repacking
overlays.
"""
def __call__(self, *args, **kwargs):
out_filename = args[0]
in_dir = args[1]
in_path = lambda *s: os.path.join(in_dir, *s)
card = NDSCard()
files_start_offset = 0
with open(in_path('header.bin'), 'rb') as fp:
card.header = NDSCardHeader(fp.read())
files_start_offset += card.header.header_size
with open(in_path('arm9.bin'), 'rb') as fp:
card.main_rom = fp.read()
files_start_offset += len(card.main_rom)
try:
with open(in_path('static_footer.bin'), 'rb') as fp:
card.static_footer = NitroFooter(fp.read())
files_start_offset += len(card.static_footer)
except FileNotFoundError:
pass
files_start_offset = align(files_start_offset, 0x200)
with open(in_path('arm7.bin'), 'rb') as fp:
card.sub_rom = fp.read()
files_start_offset += len(card.sub_rom)
files_start_offset = align(files_start_offset, 0x200)
try:
with open(in_path('banner.bin'), 'rb') as fp:
card.banner = NDSCardBanner(fp.read())
except FileNotFoundError:
pass
try:
with open(in_path('rsa_sig.bin'), 'rb') as fp:
card.rsa_sig = fp.read()
except FileNotFoundError:
pass
# NitroFS
dir_tree = list(os.walk(in_path('data/')))
names_table = Archiver._make_names_table(dir_tree)
num_files = 0
for ent_entries in names_table:
for ent_entry in ent_entries:
files_start_offset += len(ent_entry)
if isinstance(ent_entry, EntryNameTableFileEntry):
num_files += 1
dir_table = Archiver._make_dir_table(names_table)
length = len(dir_table) * len(DirectoryTableEntry())
files_start_offset += length
files_start_offset = align(files_start_offset, 0x200)
# reserve space for FAT
length = num_files * len(FileAllocationEntry())
files_start_offset += length
files_start_offset = align(files_start_offset, 0x200)
try:
files_start_offset += len(card.banner) # banner
files_start_offset = align(files_start_offset, 0x200)
except NameError:
pass
card.fat = Archiver._make_fat(dir_tree, dir_table, files_start_offset)
card.file_data = Archiver._get_files(dir_tree)
card.fnt = NDSCardFNT()
card.fnt.directory_table = dir_table
card.fnt.entry_name_table = names_table
with open(out_filename, 'wb') as fp:
fp.write(card.encode())
@staticmethod
def _make_names_table(
dir_tree: list[tuple[str, list[str], list[str]]]
) -> NDSCardFNT.EntryNameTable:
out = NDSCardFNT.EntryNameTable()
subdir_id = 0xf001
for dir_info in dir_tree:
dir_entry = []
# process directories
for dirname in dir_info[1]:
dir_entry.append(
EntryNameTableDirectoryEntry(dirname, subdir_id)
)
subdir_id += 1
# process files
for filename in dir_info[2]:
dir_entry.append(EntryNameTableFileEntry(filename))
# done with this directory, add end entry
dir_entry.append(EntryNameTableEndOfDirectoryEntry())
out.append(dir_entry)
# processed all directories
return out
@staticmethod
def _make_dir_table(names_table: NDSCardFNT.EntryNameTable) -> NDSCardFNT.DirectoryTable:
num_dirs = len(names_table)
out = [DirectoryTableEntry() for _ in range(num_dirs)]
out[0].parent_id = num_dirs # root dir
first_file_id = 0
subdir_idx = 1
for dir_idx in range(num_dirs):
# store size of corresponding FNT names table
# in `size` custom attribute
out[dir_idx].size = 0
out[dir_idx].file_id = first_file_id
for entry in names_table[dir_idx]:
out[dir_idx].size += len(entry)
if isinstance(entry, EntryNameTableDirectoryEntry):
# assign parent_id for subdirs not yet processed
out[subdir_idx].parent_id = dir_idx | 0xf000
subdir_idx += 1
elif isinstance(entry, EntryNameTableFileEntry):
first_file_id += 1
offset = num_dirs * 8
out[0].start = offset
for dir_idx in range(1, num_dirs):
offset += out[dir_idx-1].size
out[dir_idx].start = offset
# remove added `size` attributes
for entry in out:
del entry.size
# convert to DirectoryTable
out = NDSCardFNT.DirectoryTable(b''.join(
entry.encode() for entry in out
))
return out
@staticmethod
def _make_fat(
dir_tree: list[tuple[str, list[str], list[str]]],
dir_table: NDSCardFNT.DirectoryTable,
files_start_offset: int
) -> NDSCardFAT:
num_dirs = len(dir_table)
out = NDSCardFAT()
out.extend(
FileAllocationEntry()
for dir_info in dir_tree
for file in dir_info[2]
)
offset = align(files_start_offset, 0x200)
for dirtab_idx in range(num_dirs):
file_id = dir_table[dirtab_idx].file_id
for filename in dir_tree[dirtab_idx][2]:
file_path = os.path.join(dir_tree[dirtab_idx][0], filename)
length = os.stat(file_path).st_size
out[file_id].start = offset
out[file_id].end = offset + length
offset = align(offset, 0x200)
file_id += 1
return out
@staticmethod
def _get_files(dir_tree: list[tuple[str, list[str], list[str]]]) -> list[bytes]:
file_data = []
for dir_info in dir_tree:
for filename in dir_info[2]:
with open(os.path.join(dir_info[0], filename), 'rb') as fp:
file_data.append(fp.read())
return file_data