diff --git a/src/calibre/utils/zipfile.py b/src/calibre/utils/zipfile.py index 56ae082382..c5f849a073 100644 --- a/src/calibre/utils/zipfile.py +++ b/src/calibre/utils/zipfile.py @@ -3,12 +3,14 @@ Read and write ZIP files. Modified by Kovid Goyal to support replacing files in a zip archive. """ from __future__ import with_statement +import struct, os, time, sys, shutil +import binascii, cStringIO +from contextlib import closing + from calibre.ptempfile import TemporaryDirectory from calibre import sanitize_file_name from calibre.constants import filesystem_encoding from calibre.ebooks.chardet import detect -import struct, os, time, sys, shutil -import binascii, cStringIO try: import zlib # We may need its compression method @@ -134,6 +136,16 @@ _CD64_NUMBER_ENTRIES_TOTAL = 7 _CD64_DIRECTORY_SIZE = 8 _CD64_OFFSET_START_CENTDIR = 9 +def decode_arcname(name): + if not isinstance(name, unicode): + encoding = detect(name)['encoding'] + try: + name = name.decode(encoding) + except: + name = name.decode('utf-8', 'replace') + return sanitize_file_name(name.encode(filesystem_encoding, 'replace')) + + def is_zipfile(filename): """Quickly see if file is a ZIP file by checking the magic number.""" try: @@ -1026,19 +1038,12 @@ class ZipFile: targetpath = targetpath[:-1] # don't include leading "/" from file name if present - if os.path.isabs(member.filename): - targetpath = os.path.join(targetpath, member.filename[1:]) - else: - targetpath = os.path.join(targetpath, member.filename) + fname = decode_arcname(member.filename) + if fname.startswith('/'): + fname = fname[1:] + targetpath = os.path.join(targetpath, fname) targetpath = os.path.normpath(targetpath) - if not isinstance(targetpath, unicode): - encoding = detect(targetpath)['encoding'] - try: - targetpath = targetpath.decode(encoding) - except: - targetpath = targetpath.decode('utf-8', 'replace') - targetpath = targetpath.encode(filesystem_encoding) # Create all upper directories if necessary. upperdirs = os.path.dirname(targetpath) @@ -1047,16 +1052,10 @@ class ZipFile: if upperdirs and not os.path.exists(upperdirs): os.makedirs(upperdirs) - source = self.open(member, pwd=pwd) if not os.path.exists(targetpath): # Could be a previously automatically created directory - try: - target = open(targetpath, "wb") - except IOError: - targetpath = sanitize_file_name(targetpath) - target = open(targetpath, "wb") - shutil.copyfileobj(source, target) - source.close() - target.close() + with closing(self.open(member, pwd=pwd)) as source: + with open(targetpath, 'wb') as target: + shutil.copyfileobj(source, target) return targetpath @@ -1338,18 +1337,18 @@ def safe_replace(zipstream, name, datastream): names = z.infolist() with TemporaryDirectory('_zipfile_replace') as tdir: z.extractall(path=tdir) - zipstream.seek(0) - zipstream.truncate() - z = ZipFile(zipstream, 'w') path = os.path.join(tdir, *name.split('/')) shutil.copyfileobj(datastream, open(path, 'wb')) - for info in names: - current = os.path.join(tdir, *info.filename.split('/')) - if os.path.isdir(current): - z.writestr(info.filename+'/', '', 0700) - else: - z.write(current, info.filename, compress_type=info.compress_type) - z.close() + zipstream.seek(0) + zipstream.truncate() + with closing(ZipFile(zipstream, 'w')) as z: + for info in names: + fname = decode_arcname(info.filename) + current = os.path.join(tdir, *fname.split('/')) + if os.path.isdir(current): + z.writestr(info.filename+'/', '', 0700) + else: + z.write(current, info.filename, compress_type=info.compress_type) class PyZipFile(ZipFile): """Class to create ZIP archives with Python library files and packages."""