10f1469bcdad11cf8bfe79ace33d28052418ecb48borenet#!/usr/bin/env python
20f1469bcdad11cf8bfe79ace33d28052418ecb48borenet#
30f1469bcdad11cf8bfe79ace33d28052418ecb48borenet# Copyright 2016 Google Inc.
40f1469bcdad11cf8bfe79ace33d28052418ecb48borenet#
50f1469bcdad11cf8bfe79ace33d28052418ecb48borenet# Use of this source code is governed by a BSD-style license that can be
60f1469bcdad11cf8bfe79ace33d28052418ecb48borenet# found in the LICENSE file.
70f1469bcdad11cf8bfe79ace33d28052418ecb48borenet
80f1469bcdad11cf8bfe79ace33d28052418ecb48borenet
90f1469bcdad11cf8bfe79ace33d28052418ecb48borenet"""Utilities for zipping and unzipping files."""
100f1469bcdad11cf8bfe79ace33d28052418ecb48borenet
110f1469bcdad11cf8bfe79ace33d28052418ecb48borenet
120f1469bcdad11cf8bfe79ace33d28052418ecb48borenetimport fnmatch
130f1469bcdad11cf8bfe79ace33d28052418ecb48borenetimport os
140f1469bcdad11cf8bfe79ace33d28052418ecb48borenetimport zipfile
150f1469bcdad11cf8bfe79ace33d28052418ecb48borenet
160f1469bcdad11cf8bfe79ace33d28052418ecb48borenet
170f1469bcdad11cf8bfe79ace33d28052418ecb48borenetdef filtered(names, blacklist):
180f1469bcdad11cf8bfe79ace33d28052418ecb48borenet  """Filter the list of file or directory names."""
190f1469bcdad11cf8bfe79ace33d28052418ecb48borenet  rv = names[:]
200f1469bcdad11cf8bfe79ace33d28052418ecb48borenet  for pattern in blacklist:
210f1469bcdad11cf8bfe79ace33d28052418ecb48borenet    rv = [n for n in rv if not fnmatch.fnmatch(n, pattern)]
220f1469bcdad11cf8bfe79ace33d28052418ecb48borenet  return rv
230f1469bcdad11cf8bfe79ace33d28052418ecb48borenet
240f1469bcdad11cf8bfe79ace33d28052418ecb48borenet
250f1469bcdad11cf8bfe79ace33d28052418ecb48borenetdef zip(target_dir, zip_file, blacklist=None):  # pylint: disable=W0622
260f1469bcdad11cf8bfe79ace33d28052418ecb48borenet  """Zip the given directory, write to the given zip file."""
270f1469bcdad11cf8bfe79ace33d28052418ecb48borenet  if not os.path.isdir(target_dir):
280f1469bcdad11cf8bfe79ace33d28052418ecb48borenet    raise IOError('%s does not exist!' % target_dir)
290f1469bcdad11cf8bfe79ace33d28052418ecb48borenet  blacklist = blacklist or []
305580c6913500ccd80fb03a72e6992ee0d5b42d88mtklein  with zipfile.ZipFile(zip_file, 'w', zipfile.ZIP_DEFLATED, True) as z:
310f1469bcdad11cf8bfe79ace33d28052418ecb48borenet    for r, d, f in os.walk(target_dir, topdown=True):
320f1469bcdad11cf8bfe79ace33d28052418ecb48borenet      d[:] = filtered(d, blacklist)
330f1469bcdad11cf8bfe79ace33d28052418ecb48borenet      for filename in filtered(f, blacklist):
340f1469bcdad11cf8bfe79ace33d28052418ecb48borenet        filepath = os.path.join(r, filename)
350f1469bcdad11cf8bfe79ace33d28052418ecb48borenet        zi = zipfile.ZipInfo(filepath)
360f1469bcdad11cf8bfe79ace33d28052418ecb48borenet        zi.filename = os.path.relpath(filepath, target_dir)
370f1469bcdad11cf8bfe79ace33d28052418ecb48borenet        perms = os.stat(filepath).st_mode
380f1469bcdad11cf8bfe79ace33d28052418ecb48borenet        zi.external_attr = perms << 16L
395580c6913500ccd80fb03a72e6992ee0d5b42d88mtklein        zi.compress_type = zipfile.ZIP_DEFLATED
400f1469bcdad11cf8bfe79ace33d28052418ecb48borenet        with open(filepath, 'rb') as f:
410f1469bcdad11cf8bfe79ace33d28052418ecb48borenet          content = f.read()
420f1469bcdad11cf8bfe79ace33d28052418ecb48borenet        z.writestr(zi, content)
430f1469bcdad11cf8bfe79ace33d28052418ecb48borenet      for dirname in d:
440f1469bcdad11cf8bfe79ace33d28052418ecb48borenet        dirpath = os.path.join(r, dirname)
450f1469bcdad11cf8bfe79ace33d28052418ecb48borenet        z.write(dirpath, os.path.relpath(dirpath, target_dir))
460f1469bcdad11cf8bfe79ace33d28052418ecb48borenet
470f1469bcdad11cf8bfe79ace33d28052418ecb48borenet
480f1469bcdad11cf8bfe79ace33d28052418ecb48borenetdef unzip(zip_file, target_dir):
490f1469bcdad11cf8bfe79ace33d28052418ecb48borenet  """Unzip the given zip file into the target dir."""
500f1469bcdad11cf8bfe79ace33d28052418ecb48borenet  if not os.path.isdir(target_dir):
510f1469bcdad11cf8bfe79ace33d28052418ecb48borenet    os.makedirs(target_dir)
525580c6913500ccd80fb03a72e6992ee0d5b42d88mtklein  with zipfile.ZipFile(zip_file, 'r', zipfile.ZIP_DEFLATED, True) as z:
530f1469bcdad11cf8bfe79ace33d28052418ecb48borenet    for zi in z.infolist():
540f1469bcdad11cf8bfe79ace33d28052418ecb48borenet      dst_path = os.path.join(target_dir, zi.filename)
550f1469bcdad11cf8bfe79ace33d28052418ecb48borenet      if zi.filename.endswith('/'):
560f1469bcdad11cf8bfe79ace33d28052418ecb48borenet        os.mkdir(dst_path)
570f1469bcdad11cf8bfe79ace33d28052418ecb48borenet      else:
5818b84900855bb34f1694b4c50c088e7377d0e89fBrian Osman        with open(dst_path, 'wb') as f:
590f1469bcdad11cf8bfe79ace33d28052418ecb48borenet          f.write(z.read(zi))
600f1469bcdad11cf8bfe79ace33d28052418ecb48borenet      perms = zi.external_attr >> 16L
610f1469bcdad11cf8bfe79ace33d28052418ecb48borenet      os.chmod(dst_path, perms)
62