extract.py

# NOTE: This is mostly code from flame of heroesoflegend.org; I only ported it to Python 3
# Better code have been released since then; see Trails Research Group: https://github.com/Trails-Research-Group
import ISO_extract
import sys

with open(sys.argv[1], 'rb') as f:
    pathtbl = ISO_extract.getpathtbl(f)
    ISO_extract.extract_all(f, sys.argv[2], 'cp932', pathtbl)
 

filelist_builder_part1.py

# NOTE: This is mostly code from flame of heroesoflegend.org; I only ported it to Python 3
# Better code have been released since then; see Trails Research Group: https://github.com/Trails-Research-Group
import binascii
import shutil
import struct
import os

DATA_START = 0x400
buildpath = os.getcwd() + '/ISO'

def ensure_dir(dirname):
    if not os.path.exists(dirname):
        os.makedirs(dirname)
def walk_len(foldername):
    count = 0
    for dirname, dirlist, filelist in os.walk(foldername):
        count += sum(1 if len(d) < 9 else 0 for d in dirlist) + len(filelist)
    return count

print('Operation part 1: Build data.lst template file.')
ext_lst = ['']
datalst = bytearray()
for dirpath, dirlist, filelist in os.walk(buildpath):
    for filename in filelist:
        try:
            filename.encode('ascii')
        except UnicodeEncodeError:
            ensure_dir('trashed')
            shutil.copy(dirpath + '/' + filename, 'trashed')
            os.remove(dirpath + '/' + filename)
            print('{}: Moved to trashed folder, non-ascii characters '
                  'found.'.format(dirpath + '/' + filename))
for dirpath, dirlist, filelist in os.walk(buildpath):
    basepath = os.path.relpath(dirpath, start=buildpath)
    #data.lst directory entry
    if 'USRDIR' in basepath:
        if len(os.path.basename(basepath)) < 9:
            if basepath == r'PSP_GAME/USRDIR':
                datalst += b'\x00' * 8
            else:
                s = os.path.basename(basepath).encode('ascii')
                datalst += s + b'\x00' * (8 - len(s))
            datalst += struct.pack('<II', walk_len(dirpath), 0)
#Some illegal names in there. Uncomment if you want to see them.
#They are illegal (unsupported by data.lst format) so they are skipped.
##        else:
##            print(basepath)
    for filename in filelist:
        try:
            filename.encode('cp932')
        except Exception as err:
            print(filename)
            print(binascii.hexlify(filename))
            raise err
        #data.lst file entry
        if 'USRDIR' in basepath:
            file, ext = os.path.splitext(filename)
            file = file.encode('cp932')
            datalst += file + b'\x00' * (8 - len(file))
            ext = ext[1:]
            filesize = os.path.getsize(dirpath + '/' + filename)
            datalst += struct.pack('<I', filesize)
            datalst += b'\x00' * 3                  #Blank LBA entry
            if ext in ext_lst:
                datalst.append(ext_lst.index(ext))
            else:
                datalst.append(len(ext_lst))
                ext_lst.append(ext)
datalst[0x18:0x1C] = struct.pack('<I', len(datalst) + DATA_START)
with open('data.lst', 'wb') as f:
    f.write(struct.pack('<I', DATA_START + len(datalst)))
    for ext in ext_lst[1:]:
        f.write(ext.encode('ascii') + b'\00')
    f.write(b'\x00' * (DATA_START - f.tell()))
    f.write(datalst)
print('Finished building template file.',
      'Copying template into ISO folder structure.')
shutil.copy('data.lst', buildpath + r'/PSP_GAME/USRDIR/data.lst')
print('Operation part 1 finished.')

filelist_builder_part2.py

# NOTE: This is mostly code from flame of heroesoflegend.org; I only ported it to Python 3
# Better code have been released since then; see Trails Research Group: https://github.com/Trails-Research-Group
from functools import partial
from itertools import chain
import os
import struct
import ISO_extract

DATA_START = 0x400
BASEPATH = 'USRDIR'
target = 'output.iso'

print('Reading path table.')
with open(target, 'rb') as f:
    pathtbl = ISO_extract.getpathtbl(f)
print('Updating data.lst with file locations (LBA).')
datalst_size = os.path.getsize('data.lst')
with open('data.lst', 'rb+') as f:
    ext_list = []
    f.seek(4)
    for ext in iter(partial(f.read, 4), b'\x00' * 4):
        ext_list.append('.' + ext.rstrip(b'\x00').decode('ascii'))
    path_list = []
    path_count = []
    f.seek(DATA_START)
    while f.tell() < datalst_size:
        name = f.read(8).rstrip(b'\x00').decode('cp932')
        size = struct.unpack('<I', f.read(4))[0]
        LBA = f.read(3)
        ext = f.read(1)[0]
        path_count = [x - 1 for x in path_count]
        if ext == 0:
            path_list.append(name)
            path_count.append(size)
        else:
            path = 'USRDIR' + '/'.join(chain(path_list, (name + ext_list[ext - 1],)))
            LBA = ISO_extract.get_dir_rec(path, pathtbl).LBA
            f.seek(-4, 1)
            f.write(struct.pack('<I', LBA)[:3])
            f.seek(1, 1)
        while path_count[-1] == 0:
            del path_count[-1]
            del path_list[-1]
            if len(path_count) == 0:
                break
print('Inserting modified data.lst file overtop of the template.')
with open(target, 'rb+') as f:
    ISO_extract.replace(f, r'USRDIR/data.lst', 'data.lst', pathtbl)
print('Operation finished. Time to test it!')

ISO_extract.py

# NOTE: This is mostly code from flame of heroesoflegend.org; I only ported it to Python 3
# Better code have been released since then; see Trails Research Group: https://github.com/Trails-Research-Group
import pdb
import os
import struct
from collections import namedtuple

def ensure_dir(dirname):
    if not os.path.exists(dirname):
        os.makedirs(dirname)
        
def split_path(path):
    path = path.rstrip('/')
    folders = []
    while True:
        path, folder = os.path.split(path)
        if folder == '':
            if path not in ('', '/'):
                folders.append(path)
            break
        else:
            folders.append(folder)
    return folders

dir_record = namedtuple('dir_record', ['name', 'full_path', 'LBA', 'flags', 'size', 'dir_rec_pos'])
def read_dir_record(f, folder_path, codec='cp932'):
    #Advances the file position to the next directory record also.
    try:
        pos = f.tell()
        (rec_len, dummy1, LBA, dummy2, size, dummy3, dummy4, dummy5, dummy6, flags,
         dummy7, dummy8, dummy9, name_len) = (
            struct.unpack('<BBIIIIIHBBBBIB', f.read(33)))
        name = f.read(name_len)
        f.seek(pos + rec_len)
        return dir_record(name.decode(codec), folder_path + '/' + name.decode(codec),
                          LBA, flags, size, pos)
    except Exception as err:
        print(hex(pos))
        print(name)
        raise err

def getpathtbl(f, codec='cp932'):
    pathtabledata = [None]                  #Shove indexes + 1
    f.seek(16 * 0x800 + 140)                #Path table LBA
    f.seek(struct.unpack('<I', f.read(4))[0] * 0x800)
    path_record = namedtuple('path_record', ['name', 'full_path', 'LBA', 'parent_ID', 'dir_tbl'])
    while True:
        name_len, addl_len, LBA, parent_ID = struct.unpack('<BBIH', f.read(8))
        name = f.read(name_len)
        if name == b'\x00':
            name = b''
        if f.tell() % 2 != 0:
            f.read(1)
        if LBA == 0:
            break
        else:
            name_temp = name
            parent_ID_temp = parent_ID
            full_path = []
            while name_temp != b'':
                full_path.append(name_temp)
                record = pathtabledata[parent_ID_temp]
                name_temp = record.name
                parent_ID_temp = record.parent_ID
            full_path.append(b'')
            full_path = b'/'.join(reversed(full_path)).decode('ascii')
            dir_tbl = getdirtbl(f, LBA, full_path, codec)
            pathtabledata.append(path_record(name, full_path, LBA, parent_ID, dir_tbl))
    return pathtabledata

def getdirtbl(f, LBA, full_path, codec):
    saved = f.tell()
    dir_pos = LBA * 0x800
    f.seek(dir_pos)
    dir_size = read_dir_record(f, full_path, codec).size
    dir_tbl = []
    while f.tell() < dir_pos + dir_size:
        dir_rec = read_dir_record(f, full_path, codec)
        test = f.read(1)
        f.seek(-1, 1)
        #If at end of current LBA, advance to the next one
        if test == b'\x00':
            f.seek(((f.tell() + 0x800 - 1) // 0x800) * 0x800)
        if dir_rec.flags == 0:
            dir_tbl.append(dir_rec)
    f.seek(saved)
    return dir_tbl
def get_dir_rec(pathname, pathtbl):
    folder, file = os.path.split(pathname)
    dirID = getdirID(folder, pathtbl)
    if not dirID:
        return
    for dir_rec in pathtbl[dirID].dir_tbl:
        if dir_rec.name.lower() == file.lower():
            return dir_rec
    raise Exception('File {} not found in the ISO'.format(pathname))
def getdirID(pathname, pathtbl):
    folders = split_path(pathname)
    for i, record in enumerate(pathtbl[1:], 1):
        for foldername in folders:
            if record.name.decode('ascii').lower() == foldername.lower():
                prev_record = record
                record = pathtbl[record.parent_ID]
            else:
                break
        else:
            return i    #Matched full path
    else:
        return          #Not found
def extract(f, filename, LBA, size):
    f.seek(LBA * 0x800)
    with open(filename, 'wb') as g:
        g.write(f.read(size))
def extract_all(f, target_path, codec='cp932', pathtbl=None):
    if pathtbl == None:
        pathtbl = getpathtbl(f, codec)
    for path_rec in pathtbl[1:]:
        ensure_dir(target_path + path_rec.full_path)
        for dir_rec in path_rec.dir_tbl:
            with open(target_path + dir_rec.full_path, 'wb') as g:
                f.seek(dir_rec.LBA * 0x800)
                g.write(f.read(dir_rec.size))
                
def replace(f, replace_path, replace_file, pathtbl = None, LBAlist = None):
    '''replace a file in ISO f

    f must be opened for updating.
    If replace_path is a file in the ISO and replace_file is a valid file,
    the file is replaced if there's enough space for it (otherwise no action
    is taken).
    
    If pathtbl and LBAlist are not specified, it will compute them.
    If you care calling this over and over, you might want to pass them in
    to save on computation time.
    '''
    if f.mode != 'rb+':
        raise IOError('Replace: file not opened for updating')
    if pathtbl == None:
        pathtbl = getpathtbl(f)
    dir_rec = get_dir_rec(replace_path, pathtbl)
    if not os.path.isfile(replace_file):
        raise Exception('Replace file {} not found'.format(replace_file))
    if LBAlist == None:
        LBAlist = [dir_rec.LBA for path_rec in pathtbl[1:] for
                   dir_rec in path_rec.dir_tbl]
        f.seek(16 * 0x800 + 0x50)
        LBAlist.append(struct.unpack('<I', f.read(4))[0])
        LBAlist.sort()
    nextLBA = LBAlist[LBAlist.index(dir_rec.LBA) + 1]
    avail_space = (nextLBA - dir_rec.LBA) * 0x800
    size = os.path.getsize(replace_file)
    if size > avail_space:
        raise Exception('Not enough space for {} (required: {} bytes, '
                        'available: {} bytes'.format(
                            replace_file, size, avail_space))
#Time to do the replace
    f.seek(dir_rec.LBA * 0x800)         #Zero out original file
    f.write(b'\x00' * avail_space)
    f.seek(dir_rec.LBA * 0x800)
    with open(replace_file, 'rb') as g: #Write new file
        f.write(g.read())
    f.seek(dir_rec.dir_rec_pos + 10)    #Update directory record with new filesize
    f.write(struct.pack('<I', size))    
    f.write(struct.pack('>I', size))
    
def ISO_extract(filename, pathname, filelist=None):
    """Extract files from an ISO

    Doesn't necessarily have to be from a PSP ISO, any ISO should work.

    If pathname is the path to a file, filelist must not be specified.
    It will look for the specified file and if found, will extract it.

    If path name is the path to a folder, filelist must be specified.
    It will look for the folder.
    If the folder is found, it will look for each file in filelist within
    that folder and will dump any and all of them it finds.

    Pathname doesn't have to be the full path but if ambiguous, might not work.
    Only the first folder that matches the entire path will be searched.
    """
    if filelist == None:
        pathname = pathname.rstrip('/')
        pathname, targetfile = os.path.split(pathname)
        filelist = [targetfile]
    else:
        filelist = list(filelist)

    with open(filename, 'rb') as f:
        pathtbl = getpathtbl(f)
        i = getdirID(pathname, pathtbl)
        if not i:
            print('Path {} not found.'.format(pathname))
            return
        for dir_rec in pathtbl[i].dir_tbl:
            for name in filelist:
                if name.lower() == dir_rec.name.lower():
                    filelist.remove(name)
                    extract(f, dir_rec.name, dir_rec.LBA, dir_rec.size)
                    break
            if filelist == []:          #All files found
                break
    if filelist == []:                  #All files found
        pass
    else:
        print('Files not found:', ', '.join(filelist))
#Example 1:
#All of these will work for path definitions
#It does not have to be the full path
##targetfolder = r'GAMEDATA\FLD\MAP''/'
##targetfolder = r'GAMEDATA\FLD\MAP'
##targetfolder = r'\GAMEDATA\FLD\MAP' + '/'
##ISO_extract('LR.iso', targetfolder, ('MAP_T_KAN_01.BIN', 'MAP_T_KAN_00.BIN'))
#Example 2: Extract UMD_DATA.BIN from the root folder.
##ISO_extract('UMD_DATA.BIN')
#Example 3: Extract a specific pathname.
##ISO_extract(r'GAMEDATA\FLD\MAP\MAP_T_KAN_01.BIN')
##with open('LR.iso', 'rb') as f:
##    pathtbl = getpathtbl(f)
##    print(get_dir_rec(r'GAMEDATA\FLD\MAP\MAP_T_KAN_01.BIN', pathtbl))
##with open('output.iso', 'rb') as f:
##    pathtbl = getpathtbl(f)