PS2 ROMFS tool

ps2romfs.py

#!/usr/bin/env python3
import argparse
import struct
import os

ROMHeader = b"Sony Computer Entertainment Inc\x2e\0"
ConsoleTypes = {
    "C": "Retail",
    "D": "Debug",
    "T": "Developer",
    "-": "Undefined"
}
Regions = {
    "A": "America",
    "J": "Japan",
    "E": "Europe",
    "C": "China"
}

sROMDir = struct.Struct("<10sHI")

def intAlign(i, size):
    return ((i) + (size) - 1) & ~((size) - 1)

def scanROMDir(handle, reset = b"RESET\0\0\0\0\0", baseaddr = 0):
    files = {}
    offset = 0
    extOffset = 0
    special = {}
    while True:
        name, extInfoSize, fileSize = sROMDir.unpack(handle.read(sROMDir.size))
        name = name.rstrip(b"\0").decode()
        if name == "":
            break
        
        file = {
            "size": fileSize,
            "offset": offset + baseaddr,
            "extInfo": {
                "size": extInfoSize,
                "offset": extOffset
            }
        }
        
        if name not in ("RESET", "ROMDIR", "EXTINFO", "-"):
            files[name] = file
            
            ret = handle.tell()
            handle.seek(files[name]["offset"])
            if name != "ROMDIR" and handle.read(len(reset)) == reset:
                handle.seek(-len(reset), 1)
                files[name]["files"] = scanROMDir(handle, reset, files[name]["offset"])
            handle.seek(ret)
        
        else:
            special[name] = file
        
        extOffset += extInfoSize
        offset += intAlign(fileSize, 16)
    
    for file in files:
        files[file]["extInfo"]["offset"] += special["EXTINFO"]["offset"]
    
    return files

def findReset(handle, reset = b"RESET\0\0\0\0\0"):
    while True:
        for char in reset:
            test = handle.read(1)[0]
            if test == b"":
                return None
            
            if char != test:
                break
        else:
            handle.seek(-len(reset), 1)
            return handle.tell()

def findROM(d = "."):
    for file in os.scandir(d):
        #Only need files
        if not file.is_file():
            continue
        
        #PS2 BIOS is only 4 MiB
        if file.stat().st_size != 0x400000:
            continue
        
        with open(file.path, "rb") as f:
            f.seek(0x108)
            if f.read(len(ROMHeader)) == ROMHeader:
                return file.path

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Extract ps2 romfs")
    parser.add_argument("mode", choices=("l", "x", "v"),
                        help="l: List; x: Extract; v: Version")
    parser.add_argument("SCPH", default = None, nargs = "?",
                        help="Input SCPH. Default is automatic find in CWD.")
    parser.add_argument("-o", "--out", default = None,
                        help="Output directory. Default is input minus ext.")
    args = parser.parse_args()
    
    rom = None
    if args.SCPH == None:
        rom = findROM(os.getcwd())
    else:
        rom = args.SCPH
    
    
    if not rom or not os.path.isfile(rom):
        print("No ROM found.")
        parser.print_help()
        exit()
    
    if args.out == None:
        args.out = os.path.join(os.getcwd(), ".".join(os.path.basename(rom).split(".")[:1]))
    
    with open(rom, "rb") as f:
        findReset(f)
        ROMDir = scanROMDir(f)
        if args.mode == "l":
            def printdir(tree, prefix = ""):
                i = 0
                l = len(tree)
                for item in tree:
                    i += 1
                    
                    if i == l:
                        tmp = prefix + "└"
                    else:
                        tmp = prefix + "├"
                    
                    if "files" in tree[item]:
                        print(tmp + "──┬─ " + item)
                        printdir(tree[item]["files"], tmp[:-1]+"│  ")
                    else:
                        print(tmp + "──── " + item)
            
            print(os.path.basename(rom))
            printdir(ROMDir)
        
        elif args.mode == "v":
            f.seek(ROMDir["VERSTR"]["offset"])
            for t in f.read(ROMDir["VERSTR"]["size"])[2:-2].split(b"\0"):
                print(t.decode())
            
            f.seek(ROMDir["ROMVER"]["offset"])
            rv = f.read(ROMDir["ROMVER"]["size"])
            version = int(rv[0:2]) + (int(rv[2:4])/100)
            region = Regions.get(rv[4:5].decode(), "UNKNOWN REGION")
            
            ctype = ConsoleTypes.get(rv[5:6].decode(), "UNKNOWN TYPE")
            
            year = int(rv[6:10])
            month = int(rv[10:12])
            day = int(rv[12:14])
            print("Version: {:0.2f}".format(version))
            print("Region: {}".format(region))
            print("Type: {}".format(ctype))
            print("Date: {:0>4}/{:0>2}/{:0>2}".format(year, month, day))
        
        elif args.mode == "x":
            def extract(tree, dest):
                os.makedirs(dest, exist_ok = True)
                for file in tree:
                    fp = os.path.join(dest, file)
                    
                    if "files" in tree[file]:
                        extract(tree[file]["files"], fp)
                    else:
                        with open(fp, "wb") as ff:
                            f.seek(tree[file]["offset"])
                            ff.write(f.read(tree[file]["size"]))
            extract(ROMDir, args.out)
        
        else:
            print("Not implemented")