11e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)# Copyright 2013 The Chromium Authors. All rights reserved.
21e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)# Use of this source code is governed by a BSD-style license that can be
31e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)# found in the LICENSE file.
41e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)
51e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)from io import BytesIO
61e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)import posixpath
71e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)from zipfile import ZipFile
81e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)
90f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles)from compiled_file_system import SingleFile
100f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles)
111e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)
121e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)class DirectoryZipper(object):
131e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)  '''Creates zip files of whole directories.
141e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)  '''
151e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)
161e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)  def __init__(self, compiled_fs_factory, file_system):
171e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)    self._file_system = file_system
181e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)    self._zip_cache = compiled_fs_factory.Create(file_system,
191e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)                                                 self._MakeZipFile,
201e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)                                                 DirectoryZipper)
211e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)
220f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles)  # NOTE: It's ok to specify SingleFile here even though this method reads
230f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles)  # multiple files. All files are underneath |base_dir|. If any file changes its
240f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles)  # stat will change, so the stat of |base_dir| will also change.
250f1bc08d4cfcc34181b0b5cbf065c40f687bf740Torne (Richard Coles)  @SingleFile
261e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)  def _MakeZipFile(self, base_dir, files):
271e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)    base_dir = base_dir.strip('/')
281e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)    zip_bytes = BytesIO()
291e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)    with ZipFile(zip_bytes, mode='w') as zip_file:
301e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)      futures_with_name = [
311e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)          (file_name,
32a3f6a49ab37290eeeb8db0f41ec0f1cb74a68be7Torne (Richard Coles)           self._file_system.ReadSingle(posixpath.join(base_dir, file_name)))
331e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)          for file_name in files]
341e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)      for file_name, future in futures_with_name:
351e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)        file_contents = future.Get()
361e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)        if isinstance(file_contents, unicode):
371e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)          # Data is sometimes already cached as unicode.
381e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)          file_contents = file_contents.encode('utf8')
391e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)        # We want e.g. basic.zip to expand to basic/manifest.json etc, not
401e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)        # chrome/common/extensions/.../basic/manifest.json, so only use the
411e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)        # end of the path component when writing into the zip file.
421e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)        dir_name = posixpath.basename(base_dir)
431e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)        zip_file.writestr(posixpath.join(dir_name, file_name), file_contents)
441e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)    return zip_bytes.getvalue()
451e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)
461e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)  def Zip(self, path):
471e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)    '''Creates a new zip file from the recursive contents of |path| as returned
481e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)    by |_zip_cache|.
491e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)
501e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)    Paths within the zip file will be relative to and include the last
511e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)    component of |path|, such that when zipping foo/bar containing baf.txt and
521e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)    qux.txt will return a zip file which unzips to bar/baz.txt and bar/qux.txt.
531e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)    '''
541e9bf3e0803691d0a228da41fc608347b6db4340Torne (Richard Coles)    return self._zip_cache.GetFromFileListing(path)
55