tonytins/rpatool
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix some errors in Unicode and pickle handling.

Browse source
This commit is contained in:
Shiz 2017-05-09 14:31:54 +02:00
parent b484611ce4
commit f2520ad39e
1 changed files with 42 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

70
rpatool
View file

@ -10,20 +10,33 @@ import errno
import random import random
if sys.version_info[0] >= 3: if sys.version_info[0] >= 3:
def _ensure_unicode(text): def _unicode(text):
return text return text
def _prepare_unicode_for_print(text): def _printable(text):
return text return text
def _unmangle(data):
return data.encode('latin1')
def _unpickle(data):
# Specify latin1 encoding to prevent raw byte values from causing an ASCII decode error.
return pickle.loads(data, encoding='latin1')
elif sys.version_info[0] == 2: elif sys.version_info[0] == 2:
def _ensure_unicode(text): def _unicode(text):
if isinstance(text, unicode): if isinstance(text, unicode):
return text return text
return unicode(text, 'utf-8') return text.decode('utf-8')
def _prepare_unicode_for_print(text): def _printable(text):
return text.encode('utf-8') return text.encode('utf-8')
def _unmangle(data):
return data
def _unpickle(data):
return pickle.loads(data)
class RenPyArchive: class RenPyArchive:
file = None file = None
handle = None handle = None
@ -87,7 +100,8 @@ class RenPyArchive:
# Load in indexes. # Load in indexes.
self.handle.seek(offset) self.handle.seek(offset)
indexes = pickle.loads(codecs.decode(self.handle.read(), 'zlib')) contents = codecs.decode(self.handle.read(), 'zlib')
indexes = _unpickle(contents)
# Deobfuscate indexes. # Deobfuscate indexes.
if self.version == 3: if self.version == 3:
@ -131,26 +145,26 @@ class RenPyArchive:
# Check if a file exists in the archive. # Check if a file exists in the archive.
def has_file(self, filename): def has_file(self, filename):
filename = _ensure_unicode(filename) filename = _unicode(filename)
return filename in self.indexes.keys() or filename in self.files.keys() return filename in self.indexes.keys() or filename in self.files.keys()
# Read file from archive or internal storage. # Read file from archive or internal storage.
def read(self, filename): def read(self, filename):
filename = self.convert_filename(_ensure_unicode(filename)) filename = self.convert_filename(_unicode(filename))
# Check if the file exists in our indexes. # Check if the file exists in our indexes.
if filename not in self.files and filename not in self.indexes: if filename not in self.files and filename not in self.indexes:
raise IOError(errno.ENOENT, 'the requested file {0} does not exist in the given Ren\'Py archive'.format( raise IOError(errno.ENOENT, 'the requested file {0} does not exist in the given Ren\'Py archive'.format(
_prepare_unicode_for_print(filename))) _printable(filename)))
# If it's in our opened archive index, and our archive handle isn't valid, something is obviously wrong. # If it's in our opened archive index, and our archive handle isn't valid, something is obviously wrong.
if filename not in self.files and filename in self.indexes and self.handle is None: if filename not in self.files and filename in self.indexes and self.handle is None:
raise IOError(errno.ENOENT, 'the requested file {0} does not exist in the given Ren\'Py archive'.format( raise IOError(errno.ENOENT, 'the requested file {0} does not exist in the given Ren\'Py archive'.format(
_prepare_unicode_for_print(filename))) _printable(filename)))
# Check our simplified internal indexes first, in case someone wants to read a file they added before without saving, for some unholy reason. # Check our simplified internal indexes first, in case someone wants to read a file they added before without saving, for some unholy reason.
if filename in self.files: if filename in self.files:
self.verbose_print('Reading file {0} from internal storage...'.format(_prepare_unicode_for_print(filename))) self.verbose_print('Reading file {0} from internal storage...'.format(_printable(filename)))
return self.files[filename] return self.files[filename]
# We need to read the file from our open archive. # We need to read the file from our open archive.
else: else:
@ -162,13 +176,13 @@ class RenPyArchive:
prefix = '' prefix = ''
self.verbose_print('Reading file {0} from data file {1}... (offset = {2}, length = {3} bytes)'.format( self.verbose_print('Reading file {0} from data file {1}... (offset = {2}, length = {3} bytes)'.format(
_prepare_unicode_for_print(filename), self.file, offset, length)) _printable(filename), self.file, offset, length))
self.handle.seek(offset) self.handle.seek(offset)
return codecs.encode(prefix) + self.handle.read(length - len(prefix)) return _unmangle(prefix) + self.handle.read(length - len(prefix))
# Modify a file in archive or internal storage. # Modify a file in archive or internal storage.
def change(self, filename, contents): def change(self, filename, contents):
filename = _ensure_unicode(filename) filename = _unicode(filename)
# Our 'change' is basically removing the file from our indexes first, and then re-adding it. # Our 'change' is basically removing the file from our indexes first, and then re-adding it.
self.remove(filename) self.remove(filename)
@ -176,29 +190,29 @@ class RenPyArchive:
# Add a file to the internal storage. # Add a file to the internal storage.
def add(self, filename, contents): def add(self, filename, contents):
filename = self.convert_filename(_ensure_unicode(filename)) filename = self.convert_filename(_unicode(filename))
if filename in self.files or filename in self.indexes: if filename in self.files or filename in self.indexes:
raise ValueError('file {0} already exists in archive'.format(_prepare_unicode_for_print(filename))) raise ValueError('file {0} already exists in archive'.format(_printable(filename)))
self.verbose_print('Adding file {0} to archive... (length = {1} bytes)'.format( self.verbose_print('Adding file {0} to archive... (length = {1} bytes)'.format(
_prepare_unicode_for_print(filename), len(contents))) _printable(filename), len(contents)))
self.files[filename] = contents self.files[filename] = contents
# Remove a file from archive or internal storage. # Remove a file from archive or internal storage.
def remove(self, filename): def remove(self, filename):
filename = _ensure_unicode(filename) filename = _unicode(filename)
if filename in self.files: if filename in self.files:
self.verbose_print('Removing file {0} from internal storage...'.format(_prepare_unicode_for_print(filename))) self.verbose_print('Removing file {0} from internal storage...'.format(_printable(filename)))
del self.files[filename] del self.files[filename]
elif filename in self.indexes: elif filename in self.indexes:
self.verbose_print('Removing file {0} from archive indexes...'.format(_prepare_unicode_for_print(filename))) self.verbose_print('Removing file {0} from archive indexes...'.format(_printable(filename)))
del self.indexes[filename] del self.indexes[filename]
else: else:
raise IOError(errno.ENOENT, 'the requested file {0} does not exist in this archive'.format(_prepare_unicode_for_print(filename))) raise IOError(errno.ENOENT, 'the requested file {0} does not exist in this archive'.format(_printable(filename)))
# Load archive. # Load archive.
def load(self, filename): def load(self, filename):
filename = _ensure_unicode(filename) filename = _unicode(filename)
if self.handle is not None: if self.handle is not None:
self.handle.close() self.handle.close()
@ -210,7 +224,7 @@ class RenPyArchive:
# Save current state into a new file, merging archive and internal storage, rebuilding indexes, and optionally saving in another format version. # Save current state into a new file, merging archive and internal storage, rebuilding indexes, and optionally saving in another format version.
def save(self, filename = None): def save(self, filename = None):
filename = _ensure_unicode(filename) filename = _unicode(filename)
if filename is None: if filename is None:
filename = self.file filename = self.file
@ -322,17 +336,17 @@ if __name__ == "__main__":
# Determine output file/directory and input archive # Determine output file/directory and input archive
if arguments.create: if arguments.create:
archive = None archive = None
output = _ensure_unicode(arguments.archive) output = _unicode(arguments.archive)
else: else:
archive = _ensure_unicode(arguments.archive) archive = _unicode(arguments.archive)
if 'outfile' in arguments and arguments.outfile is not None: if 'outfile' in arguments and arguments.outfile is not None:
output = _ensure_unicode(arguments.outfile) output = _unicode(arguments.outfile)
else: else:
# Default output directory for extraction is the current directory. # Default output directory for extraction is the current directory.
if arguments.extract: if arguments.extract:
output = '.' output = '.'
else: else:
output = _ensure_unicode(arguments.archive) output = _unicode(arguments.archive)
# Normalize files. # Normalize files.
if len(arguments.files) > 0 and isinstance(arguments.files[0], list): if len(arguments.files) > 0 and isinstance(arguments.files[0], list):
@ -367,7 +381,7 @@ if __name__ == "__main__":
# Iterate over the given files to add to archive. # Iterate over the given files to add to archive.
for filename in arguments.files: for filename in arguments.files:
add_file(_ensure_unicode(filename)) add_file(_unicode(filename))
# Set version for saving, and save. # Set version for saving, and save.
archive.version = version archive.version = version