rename_vpks_threaded.py

# -*- coding: utf-8 -*-
# Original by reddit.com/u/nomansuniverse.
# Threaded version by yours truly.

from __future__ import print_function
from multiprocessing.pool import ThreadPool

import argparse
import filecmp
import glob
import os
import string
import struct
import sys
import zipfile

def rename_file((filename, args)):
    #print("%s %s" % (filename, args))
    filename = filename.strip()
    print("Renaming '%s'...\n" % filename, end='')
    
    if not os.path.exists(filename):
        print("Could not find %s, skipping...\n" % filename, end='')
        return
    
    # Indenting got hard to read so with removed with constructor
    # except moved moved here
    try:
        vpk = zipfile.ZipFile(filename, 'r')
    except zipfile.BadZipfile:
        # i get FAILED_FAILED if running multiple times.
        if args.invalidcheck and 'FAILED_' not in filename:
            # If it's not a valid ZIP file (thus, not a valid VPK), mark it as such
            failed_filename = "FAILED_%s" % filename
            os.rename(filename, failed_filename)
        return
    
    param_path = "sce_sys/param.sfo"
    if param_path not in vpk.namelist():
        param_path = None
        for subfile in vpk.namelist(): # Search for param.sfo
            if subfile.endswith("param.sfo"):
                param_path = subfile
                break
            
    if not param_path:
        print("Could not find param.sfo in %s, skipping...\n" % filename, end='')
        return
        
    # Read param.sfo
    param_data = vpk.read(param_path)
    
    key_offset, data_offset = struct.unpack("<II", param_data[0x08:0x10])
    keys = param_data[key_offset:data_offset].strip('\0').split('\0')
      
    def get_data(key):
        if key not in keys:
            return None
    
        info_offset = 0x14 + (keys.index(key) * 0x10)
        length, _, offset = struct.unpack("<III", param_data[info_offset + 0x04: info_offset + 0x10])                    
        return param_data[data_offset+offset:data_offset+offset+length]
        
    title = get_data("TITLE").strip('\0').strip().decode('utf-8')
    title_id = get_data("TITLE_ID").strip('\0').strip()
    version = get_data("APP_VER").strip('\0').strip()
        
    if not title:
        print("Could not find title of game for %s\n" % filename, end='')
        exit(-1)
        
    region = None
    if title_id in  ["PCSB", "PCSF"]:
        region = "EUR"
    elif title_id in ["PCSA", "PCSE"]:
        region = "USA"
    elif title_id in  ["PCSG", "PCSC"]:
        region = "JPN"
    else:
        region = "UNK"
        
    new_filename_base = "%s" % title
    if title_id:
        new_filename_base = "%s [%s]" % (new_filename_base, title_id)
    if version:
        new_filename_base = "%s (v%s)" % (new_filename_base, version)
    if region:
        new_filename_base = "%s (%s)" % (new_filename_base, region)
    
    # Remove invalid characters from filename in the ASCII range
    invalid_chars = u'™©®"<>|:*?\\/'
    new_filename_base = ''.join('_' if (ord(c) > 0x7f and args.asciionly) else c for c in new_filename_base if (c not in invalid_chars and ord(c) >= 0x20))
    new_filename = "%s.%s" % (new_filename_base, args.ext)
    
    # Check if the files are the same
    # A simple filename comparison does not work because of encoding issues unfortunately
    # And such a check is required for when 2 VPKs are required because they will return the same name
    dupe = 1
    while os.path.exists(new_filename) and not filecmp.cmp(filename, new_filename):
        new_filename = "%s (%d).%s" % (new_filename_base, dupe, args.ext)
        dupe += 1
    
    new_filename = new_filename.encode(args.charset, errors='ignore') # Changes things like superscripts
    
    # no point renaming if same, helps running multiple times
    if filename != new_filename:
        os.rename(filename, new_filename)
    
    vpk.close()

def find_files(args):
    # Threadpool for 
    pool = ThreadPool()
    pool.map(rename_file, [(f, args) for f in glob.glob(args.input)])
    pool.close()
    pool.join()

def main():
    parser = argparse.ArgumentParser(description='Rename VPKs')
    parser.add_argument('--input', '-i', default="*.vpk",
                        help='Input filename (wildcard acceptable)')
    parser.add_argument('--ext', '-e', default="vpk",
                        help='Output extension type. Default: vpk')
    parser.add_argument('--asciionly', '-a', action='store_true', default=False,
                        help='Output ASCII-restricted filenames')
    parser.add_argument('--invalidcheck', '-f', action='store_true', default=True,
                        help='Rename invalid archives')
    parser.add_argument('--charset', '-c', default=sys.getfilesystemencoding(),
                        help='Output filename charset. Default: system default charset')
    args = parser.parse_args()
    
    find_files(args)
        
if __name__ == '__main__':
    main()