Add asset management scripts

These provide an easy way to create assets to be used by bots,
eg. Android SDK.

To create an asset:
$ infra/bots/assets/assets.py add android_sdk
(adds scripts in infra/bots/assets/android_sdk)

To upload a new version of an asset:
$ infra/bots/assets/android_sdk/upload.py -t $ANDROID_SDK_ROOT
(uploads Android SDK to GS, writes a version file)
$ git commit
$ git cl upload

To download the current version of the asset:
$ infra/bots/assets/android_sdk/download.py -t ../tmp

BUG=skia:5427
GOLD_TRYBOT_URL= https://gold.skia.org/search?issue=2069543002

Review-Url: https://codereview.chromium.org/2069543002

14 files changed, 777 insertions, 0 deletions

diff --git a/infra/bots/assets/README.md b/infra/bots/assets/README.md
new file mode 100644
index 0000000000..91d0247983
--- /dev/null
+++ b/infra/bots/assets/README.md
@@ -0,0 +1,47 @@
+Assets
+======
+
+This directory contains tooling for managing assets used by the bots.  The
+primary entry point is assets.py, which allows a user to add, remove, upload,
+and download assets.
+
+Assets are stored in Google Storage, named for their version number.
+
+
+Individual Assets
+-----------------
+
+Each asset has its own subdirectory with the following contents:
+* VERSION:  The current version number of the asset.
+* download.py:  Convenience script for downloading the current version of the asset.
+* upload.py:  Convenience script for uploading a new version of the asset.
+* [optional] create.py:  Script which creates the asset, implemented by the user.
+* [optional] create\_and\_upload.py:  Convenience script which combines create.py with upload.py.
+
+
+Examples
+-------
+
+Add a new asset and upload an initial version.
+
+```
+$ infra/bots/assets/assets.py add myasset
+Creating asset in infra/bots/assets/myasset
+Creating infra/bots/assets/myasset/download.py
+Creating infra/bots/assets/myasset/upload.py
+Creating infra/bots/assets/myasset/common.py
+Add script to automate creation of this asset? (y/n) n
+$ infra/bots/assets/myasset/upload.py -t ${MY_ASSET_LOCATION}
+$ git commit
+```
+
+Add an asset whose creation can be automated.
+
+```
+$ infra/bots/assets/assets.py add myasset
+Add script to automate creation of this asset? (y/n) y
+$ vi infra/bots/assets/myasset/create.py
+(implement the create_asset function)
+$ infra/bots/assets/myasset/create_and_upload.py
+$ git commit
+```
diff --git a/infra/bots/assets/__init__.py b/infra/bots/assets/__init__.py
new file mode 100644
index 0000000000..78953f5f51
--- /dev/null
+++ b/infra/bots/assets/__init__.py
@@ -0,0 +1,6 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
diff --git a/infra/bots/assets/asset_utils.py b/infra/bots/assets/asset_utils.py
new file mode 100644
index 0000000000..a13d2290a3
--- /dev/null
+++ b/infra/bots/assets/asset_utils.py
@@ -0,0 +1,174 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Utilities for managing assets."""
+
+
+import argparse
+import os
+import shlex
+import shutil
+import subprocess
+import sys
+
+SKIA_DIR = os.path.abspath(os.path.realpath(os.path.join(
+    os.path.dirname(os.path.abspath(__file__)),
+    os.pardir, os.pardir, os.pardir)))
+INFRA_BOTS_DIR = os.path.join(SKIA_DIR, 'infra', 'bots')
+sys.path.insert(0, INFRA_BOTS_DIR)
+import utils
+import zip_utils
+
+
+ASSETS_DIR = os.path.join(INFRA_BOTS_DIR, 'assets')
+DEFAULT_GS_BUCKET = 'skia-buildbots'
+GS_SUBDIR_TMPL = 'gs://%s/assets/%s'
+GS_PATH_TMPL = '%s/%s.zip'
+VERSION_FILENAME = 'VERSION'
+ZIP_BLACKLIST = ['.git', '.svn', '*.pyc', '.DS_STORE']
+
+
+class _GSWrapper(object):
+  """Wrapper object for interacting with Google Storage."""
+  def __init__(self, gsutil):
+    gsutil = os.path.abspath(gsutil) if gsutil else 'gsutil'
+    self._gsutil = [gsutil]
+    if gsutil.endswith('.py'):
+      self._gsutil = ['python', gsutil]
+
+  def copy(self, src, dst):
+    """Copy src to dst."""
+    subprocess.check_call(self._gsutil + ['cp', src, dst])
+
+  def list(self, path):
+    """List objects in the given path."""
+    try:
+      return subprocess.check_output(self._gsutil + ['ls', path]).splitlines()
+    except subprocess.CalledProcessError:
+      # If the prefix does not exist, we'll get an error, which is okay.
+      return []
+
+
+def _prompt(prompt):
+  """Prompt for input, return result."""
+  return raw_input(prompt)
+
+
+class Asset(object):
+  def __init__(self, name, gs_bucket=DEFAULT_GS_BUCKET, gsutil=None):
+    self._gs = _GSWrapper(gsutil)
+    self._gs_subdir = GS_SUBDIR_TMPL % (gs_bucket, name)
+    self._name = name
+    self._dir = os.path.join(ASSETS_DIR, self._name)
+
+  @property
+  def version_file(self):
+    """Return the path to the version file for this asset."""
+    return os.path.join(self._dir, VERSION_FILENAME)
+
+  def get_current_version(self):
+    """Obtain the current version of the asset."""
+    if not os.path.isfile(self.version_file):
+      return -1
+    with open(self.version_file) as f:
+      return int(f.read())
+
+  def get_available_versions(self):
+    """Return the existing version numbers for this asset."""
+    files = self._gs.list(self._gs_subdir)
+    bnames = [os.path.basename(f) for f in files]
+    suffix = '.zip'
+    versions = [int(f[:-len(suffix)]) for f in bnames if f.endswith(suffix)]
+    versions.sort()
+    return versions
+
+  def get_next_version(self):
+    """Find the next available version number for the asset."""
+    versions = self.get_available_versions()
+    if len(versions) == 0:
+      return 0
+    return versions[-1] + 1
+
+  def download_version(self, version, target_dir):
+    """Download the specified version of the asset."""
+    gs_path = GS_PATH_TMPL % (self._gs_subdir, str(version))
+    target_dir = os.path.abspath(target_dir)
+    with utils.tmp_dir():
+      zip_file = os.path.join(os.getcwd(), '%d.zip' % version)
+      self._gs.copy(gs_path, zip_file)
+      zip_utils.unzip(zip_file, target_dir)
+
+  def download_current_version(self, target_dir):
+    """Download the version of the asset specified in its version file."""
+    v = self.get_current_version()
+    self.download_version(v, target_dir)
+
+  def upload_new_version(self, target_dir, commit=False):
+    """Upload a new version and update the version file for the asset."""
+    version = self.get_next_version()
+    target_dir = os.path.abspath(target_dir)
+    with utils.tmp_dir():
+      zip_file = os.path.join(os.getcwd(), '%d.zip' % version)
+      zip_utils.zip(target_dir, zip_file, blacklist=ZIP_BLACKLIST)
+      gs_path = GS_PATH_TMPL % (self._gs_subdir, str(version))
+      self._gs.copy(zip_file, gs_path)
+
+    def _write_version():
+      with open(self.version_file, 'w') as f:
+        f.write(str(version))
+      subprocess.check_call([utils.GIT, 'add', self.version_file])
+
+    with utils.chdir(SKIA_DIR):
+      if commit:
+        with utils.git_branch():
+          _write_version()
+          subprocess.check_call([
+              utils.GIT, 'commit', '-m', 'Update %s version' % self._name])
+          subprocess.check_call([utils.GIT, 'cl', 'upload', '--bypass-hooks'])
+      else:
+        _write_version()
+
+  @classmethod
+  def add(cls, name, gs_bucket=DEFAULT_GS_BUCKET, gsutil=None):
+    """Add an asset."""
+    asset = cls(name, gs_bucket=gs_bucket, gsutil=gsutil)
+    if os.path.isdir(asset._dir):
+      raise Exception('Asset %s already exists!' % asset._name)
+
+    print 'Creating asset in %s' % asset._dir
+    os.mkdir(asset._dir)
+    def copy_script(script):
+      src = os.path.join(ASSETS_DIR, 'scripts', script)
+      dst = os.path.join(asset._dir, script)
+      print 'Creating %s' % dst
+      shutil.copy(src, dst)
+      subprocess.check_call([utils.GIT, 'add', dst])
+
+    for script in ('download.py', 'upload.py', 'common.py'):
+      copy_script(script)
+    resp = _prompt('Add script to automate creation of this asset? (y/n) ')
+    if resp == 'y':
+      copy_script('create.py')
+      copy_script('create_and_upload.py')
+      print 'You will need to add implementation to the creation script.'
+    print 'Successfully created asset %s.' % asset._name
+    return asset
+
+  def remove(self):
+    """Remove this asset."""
+    # Ensure that the asset exists.
+    if not os.path.isdir(self._dir):
+      raise Exception('Asset %s does not exist!' % self._name)
+
+    # Remove the asset.
+    subprocess.check_call([utils.GIT, 'rm', '-rf', self._dir])
+    if os.path.isdir(self._dir):
+      shutil.rmtree(self._dir)
+
+    # We *could* remove all uploaded versions of the asset in Google Storage but
+    # we choose not to be that destructive.
diff --git a/infra/bots/assets/asset_utils_test.py b/infra/bots/assets/asset_utils_test.py
new file mode 100644
index 0000000000..edfc271f9e
--- /dev/null
+++ b/infra/bots/assets/asset_utils_test.py
@@ -0,0 +1,124 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Tests for asset_utils."""
+
+
+import asset_utils
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+import unittest
+import uuid
+
+
+FILE_DIR = os.path.dirname(os.path.abspath(__file__))
+INFRA_BOTS_DIR = os.path.realpath(os.path.join(
+    FILE_DIR, os.pardir, 'infra', 'bots'))
+sys.path.insert(0, INFRA_BOTS_DIR)
+import test_utils
+import utils
+
+
+GS_BUCKET = 'skia-infra-testdata'
+
+
+def _fake_prompt(result):
+  """Make a function that pretends to prompt for input and returns a result."""
+  return lambda s: result
+
+
+def _write_stuff(target_dir):
+  """Write some files and directories into target_dir."""
+  fw = test_utils.FileWriter(target_dir)
+  fw.mkdir('mydir')
+  fw.mkdir('anotherdir', 0666)
+  fw.mkdir('dir3', 0600)
+  fw.mkdir('subdir')
+  fw.write('a.txt', 0777)
+  fw.write('b.txt', 0751)
+  fw.write('c.txt', 0640)
+  fw.write(os.path.join('subdir', 'd.txt'), 0640)
+
+
+class AssetUtilsTest(unittest.TestCase):
+  def setUp(self):
+    self.asset_name = str(uuid.uuid4())
+    self.old_prompt = asset_utils._prompt
+    asset_utils._prompt = _fake_prompt('y')
+    self.a = asset_utils.Asset.add(self.asset_name, gs_bucket=GS_BUCKET)
+
+  def tearDown(self):
+    if self.a:
+      self.a.remove()
+    asset_utils._prompt = self.old_prompt
+
+    gs_path = 'gs://%s/assets/%s' % (GS_BUCKET, self.asset_name)
+    attempt_delete = True
+    try:
+      subprocess.check_call(['gsutil', 'ls', gs_path])
+    except subprocess.CalledProcessError:
+      attempt_delete = False
+    if attempt_delete:
+      subprocess.check_call(['gsutil', 'rm', '-rf', gs_path])
+
+  def test_add_remove(self):
+    # Ensure that we can't create an asset twice.
+    with self.assertRaises(Exception):
+      asset_utils.Asset.add(self.asset_name, gs_bucket=GS_BUCKET)
+
+    # Ensure that the asset dir exists.
+    asset_dir = os.path.join(FILE_DIR, self.asset_name)
+    self.assertTrue(os.path.isdir(asset_dir))
+
+    # Remove the asset, ensure that it's gone.
+    self.a.remove()
+    self.a = None
+    self.assertFalse(os.path.exists(asset_dir))
+
+  def test_upload_download(self):
+    with utils.tmp_dir():
+      # Create input files and directories.
+      input_dir = os.path.join(os.getcwd(), 'input')
+      _write_stuff(input_dir)
+
+      # Upload a version, download it again.
+      self.a.upload_new_version(input_dir)
+      output_dir = os.path.join(os.getcwd(), 'output')
+      self.a.download_current_version(output_dir)
+
+      # Compare.
+      test_utils.compare_trees(self, input_dir, output_dir)
+
+  def test_versions(self):
+    with utils.tmp_dir():
+      # Create input files and directories.
+      input_dir = os.path.join(os.getcwd(), 'input')
+      _write_stuff(input_dir)
+
+      self.assertEqual(self.a.get_current_version(), -1)
+      self.assertEqual(self.a.get_available_versions(), [])
+      self.assertEqual(self.a.get_next_version(), 0)
+
+      self.a.upload_new_version(input_dir)
+
+      self.assertEqual(self.a.get_current_version(), 0)
+      self.assertEqual(self.a.get_available_versions(), [0])
+      self.assertEqual(self.a.get_next_version(), 1)
+
+      self.a.upload_new_version(input_dir)
+
+      self.assertEqual(self.a.get_current_version(), 1)
+      self.assertEqual(self.a.get_available_versions(), [0, 1])
+      self.assertEqual(self.a.get_next_version(), 2)
+
+
+if __name__ == '__main__':
+  unittest.main()
diff --git a/infra/bots/assets/assets.py b/infra/bots/assets/assets.py
new file mode 100755
index 0000000000..538a41b585
--- /dev/null
+++ b/infra/bots/assets/assets.py
@@ -0,0 +1,80 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Tool for managing assets."""
+
+
+import argparse
+import asset_utils
+import os
+import shutil
+import subprocess
+import sys
+
+FILE_DIR = os.path.dirname(os.path.abspath(__file__))
+INFRA_BOTS_DIR = os.path.realpath(os.path.join(FILE_DIR, os.pardir))
+
+sys.path.insert(0, INFRA_BOTS_DIR)
+import utils
+
+
+def add(args):
+  """Add a new asset."""
+  asset_utils.Asset.add(args.asset_name)
+
+
+def remove(args):
+  """Remove an asset."""
+  asset_utils.Asset(args.asset_name).remove()
+
+
+def download(args):
+  """Download the current version of an asset."""
+  asset = asset_utils.Asset(args.asset_name, gsutil=args.gsutil)
+  asset.download_current_version(args.target_dir)
+
+
+def upload(args):
+  """Upload a new version of the asset."""
+  asset = asset_utils.Asset(args.asset_name, gsutil=args.gsutil)
+  asset.upload_new_version(args.target_dir, commit=args.commit)
+
+
+def main(argv):
+  parser = argparse.ArgumentParser(description='Tool for managing assets.')
+  subs = parser.add_subparsers(help='Commands:')
+
+  prs_add = subs.add_parser('add', help='Add a new asset.')
+  prs_add.set_defaults(func=add)
+  prs_add.add_argument('asset_name', help='Name of the asset.')
+
+  prs_remove = subs.add_parser('remove', help='Remove an asset.')
+  prs_remove.set_defaults(func=remove)
+  prs_remove.add_argument('asset_name', help='Name of the asset.')
+
+  prs_download = subs.add_parser(
+      'download', help='Download the current version of an asset.')
+  prs_download.set_defaults(func=download)
+  prs_download.add_argument('asset_name', help='Name of the asset.')
+  prs_download.add_argument('--target_dir', '-t', required=True)
+  prs_download.add_argument('--gsutil')
+
+  prs_upload = subs.add_parser(
+      'upload', help='Upload a new version of an asset.')
+  prs_upload.set_defaults(func=upload)
+  prs_upload.add_argument('asset_name', help='Name of the asset.')
+  prs_upload.add_argument('--target_dir', '-t', required=True)
+  prs_upload.add_argument('--gsutil')
+  prs_upload.add_argument('--commit', action='store_true')
+
+  args = parser.parse_args(argv)
+  args.func(args)
+
+
+if __name__ == '__main__':
+  main(sys.argv[1:])
diff --git a/infra/bots/assets/scripts/common.py b/infra/bots/assets/scripts/common.py
new file mode 100755
index 0000000000..4920c9b4fb
--- /dev/null
+++ b/infra/bots/assets/scripts/common.py
@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Common vars used by scripts in this directory."""
+
+
+import os
+import sys
+
+FILE_DIR = os.path.dirname(os.path.abspath(__file__))
+INFRA_BOTS_DIR = os.path.realpath(os.path.join(FILE_DIR, os.pardir, os.pardir))
+
+sys.path.insert(0, INFRA_BOTS_DIR)
+from assets import assets
+
+ASSET_NAME = os.path.basename(FILE_DIR)
+
+
+def run(cmd):
+  """Run a command, eg. "upload" or "download". """
+  assets.main([cmd, ASSET_NAME] + sys.argv[1:])
diff --git a/infra/bots/assets/scripts/create.py b/infra/bots/assets/scripts/create.py
new file mode 100755
index 0000000000..4f176085fb
--- /dev/null
+++ b/infra/bots/assets/scripts/create.py
@@ -0,0 +1,28 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Create the asset."""
+
+
+import argparse
+
+
+def create_asset(target_dir):
+  """Create the asset."""
+  raise NotImplementedError('Implement me!')
+
+
+def main():
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--target_dir', '-t', required=True)
+  args = parser.parse_args()
+  create_asset(args.target_dir)
+
+
+if __name__ == '__main__':
+  main()
diff --git a/infra/bots/assets/scripts/create_and_upload.py b/infra/bots/assets/scripts/create_and_upload.py
new file mode 100755
index 0000000000..1356447477
--- /dev/null
+++ b/infra/bots/assets/scripts/create_and_upload.py
@@ -0,0 +1,42 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Create the asset and upload it."""
+
+
+import argparse
+import common
+import os
+import subprocess
+import sys
+import utils
+
+
+def main():
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--gsutil')
+  args = parser.parse_args()
+
+  with utils.tmp_dir():
+    cwd = os.getcwd()
+    create_script = os.path.join(common.FILE_DIR, 'create.py')
+    upload_script = os.path.join(common.FILE_DIR, 'upload.py')
+
+    try:
+      subprocess.check_call(['python', create_script, '-t', cwd])
+      cmd = ['python', upload_script, '-t', cwd]
+      if args.gsutil:
+        cmd.extend(['--gsutil', args.gsutil])
+      subprocess.check_call(cmd)
+    except subprocess.CalledProcessError:
+      # Trap exceptions to avoid printing two stacktraces.
+      sys.exit(1)
+
+
+if __name__ == '__main__':
+  main()
diff --git a/infra/bots/assets/scripts/download.py b/infra/bots/assets/scripts/download.py
new file mode 100755
index 0000000000..96cc87d43f
--- /dev/null
+++ b/infra/bots/assets/scripts/download.py
@@ -0,0 +1,16 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Download the current version of the asset."""
+
+
+import common
+
+
+if __name__ == '__main__':
+  common.run('download')
diff --git a/infra/bots/assets/scripts/upload.py b/infra/bots/assets/scripts/upload.py
new file mode 100755
index 0000000000..ba7fc8b6a1
--- /dev/null
+++ b/infra/bots/assets/scripts/upload.py
@@ -0,0 +1,16 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Upload a new version of the asset."""
+
+
+import common
+
+
+if __name__ == '__main__':
+  common.run('upload')
diff --git a/infra/bots/download_asset.isolate b/infra/bots/download_asset.isolate
new file mode 100644
index 0000000000..0c4d8738bc
--- /dev/null
+++ b/infra/bots/download_asset.isolate
@@ -0,0 +1,10 @@
+{
+  'includes': [
+    'infrabots.isolate',
+  ],
+  'variables': {
+    'command': [
+      'python', 'assets/<(ASSET)/download.py', '-t', '${ISOLATED_OUTDIR}', '--gsutil', '<(GSUTIL)',
+    ],
+  },
+}
diff --git a/infra/bots/test_utils.py b/infra/bots/test_utils.py
new file mode 100644
index 0000000000..aa7ac0fb7e
--- /dev/null
+++ b/infra/bots/test_utils.py
@@ -0,0 +1,73 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Test utilities."""
+
+
+import filecmp
+import os
+import uuid
+
+
+class FileWriter(object):
+  """Write files into a given directory."""
+  def __init__(self, cwd):
+    self._cwd = cwd
+    if not os.path.exists(self._cwd):
+      os.makedirs(self._cwd)
+
+  def mkdir(self, dname, mode=0755):
+    """Create the given directory with the given mode."""
+    dname = os.path.join(self._cwd, dname)
+    os.mkdir(dname)
+    os.chmod(dname, mode)
+
+  def write(self, fname, mode=0640):
+    """Write the file with the given mode and random contents."""
+    fname = os.path.join(self._cwd, fname)
+    with open(fname, 'w') as f:
+      f.write(str(uuid.uuid4()))
+    os.chmod(fname, mode)
+
+  def remove(self, fname):
+    """Remove the file."""
+    fname = os.path.join(self._cwd, fname)
+    if os.path.isfile(fname):
+      os.remove(fname)
+    else:
+      os.rmdir(fname)
+
+
+def compare_trees(test, a, b):
+  """Compare two directory trees, assert if any differences."""
+  def _cmp(prefix, dcmp):
+    # Verify that the file and directory names are the same.
+    test.assertEqual(len(dcmp.left_only), 0)
+    test.assertEqual(len(dcmp.right_only), 0)
+    test.assertEqual(len(dcmp.diff_files), 0)
+    test.assertEqual(len(dcmp.funny_files), 0)
+
+    # Verify that the files are identical.
+    for f in dcmp.common_files:
+      pathA = os.path.join(a, prefix, f)
+      pathB = os.path.join(b, prefix, f)
+      test.assertTrue(filecmp.cmp(pathA, pathB, shallow=False))
+      statA = os.stat(pathA)
+      statB = os.stat(pathB)
+      test.assertEqual(statA.st_mode, statB.st_mode)
+      with open(pathA, 'rb') as f:
+        contentsA = f.read()
+      with open(pathB, 'rb') as f:
+        contentsB = f.read()
+      test.assertEqual(contentsA, contentsB)
+
+    # Recurse on subdirectories.
+    for prefix, obj in dcmp.subdirs.iteritems():
+      _cmp(prefix, obj)
+
+  _cmp('', filecmp.dircmp(a, b))
diff --git a/infra/bots/zip_utils.py b/infra/bots/zip_utils.py
new file mode 100644
index 0000000000..7f269b9e9b
--- /dev/null
+++ b/infra/bots/zip_utils.py
@@ -0,0 +1,61 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Utilities for zipping and unzipping files."""
+
+
+import fnmatch
+import os
+import zipfile
+
+
+def filtered(names, blacklist):
+  """Filter the list of file or directory names."""
+  rv = names[:]
+  for pattern in blacklist:
+    rv = [n for n in rv if not fnmatch.fnmatch(n, pattern)]
+  return rv
+
+
+def zip(target_dir, zip_file, blacklist=None):  # pylint: disable=W0622
+  """Zip the given directory, write to the given zip file."""
+  if not os.path.isdir(target_dir):
+    raise IOError('%s does not exist!' % target_dir)
+  blacklist = blacklist or []
+  with zipfile.ZipFile(zip_file, 'w') as z:
+    for r, d, f in os.walk(target_dir, topdown=True):
+      d[:] = filtered(d, blacklist)
+      for filename in filtered(f, blacklist):
+        filepath = os.path.join(r, filename)
+        zi = zipfile.ZipInfo(filepath)
+        zi.filename = os.path.relpath(filepath, target_dir)
+        perms = os.stat(filepath).st_mode
+        zi.external_attr = perms << 16L
+        zi.compress_type = zipfile.ZIP_STORED
+        with open(filepath, 'rb') as f:
+          content = f.read()
+        z.writestr(zi, content)
+      for dirname in d:
+        dirpath = os.path.join(r, dirname)
+        z.write(dirpath, os.path.relpath(dirpath, target_dir))
+
+
+def unzip(zip_file, target_dir):
+  """Unzip the given zip file into the target dir."""
+  if not os.path.isdir(target_dir):
+    os.makedirs(target_dir)
+  with zipfile.ZipFile(zip_file, 'r') as z:
+    for zi in z.infolist():
+      dst_path = os.path.join(target_dir, zi.filename)
+      if zi.filename.endswith('/'):
+        os.mkdir(dst_path)
+      else:
+        with open(dst_path, 'w') as f:
+          f.write(z.read(zi))
+      perms = zi.external_attr >> 16L
+      os.chmod(dst_path, perms)
diff --git a/infra/bots/zip_utils_test.py b/infra/bots/zip_utils_test.py
new file mode 100644
index 0000000000..4f88a115df
--- /dev/null
+++ b/infra/bots/zip_utils_test.py
@@ -0,0 +1,74 @@
+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Tests for zip_utils."""
+
+
+import filecmp
+import os
+import test_utils
+import unittest
+import utils
+import uuid
+import zip_utils
+
+
+class ZipUtilsTest(unittest.TestCase):
+  def test_zip_unzip(self):
+    with utils.tmp_dir():
+      fw = test_utils.FileWriter(os.path.join(os.getcwd(), 'input'))
+      # Create input files and directories.
+      fw.mkdir('mydir')
+      fw.mkdir('anotherdir', 0666)
+      fw.mkdir('dir3', 0600)
+      fw.mkdir('subdir')
+      fw.write('a.txt', 0777)
+      fw.write('b.txt', 0751)
+      fw.write('c.txt', 0640)
+      fw.write(os.path.join('subdir', 'd.txt'), 0640)
+      
+      # Zip, unzip.
+      zip_utils.zip('input', 'test.zip')
+      zip_utils.unzip('test.zip', 'output')
+
+      # Compare the inputs and outputs.
+      test_utils.compare_trees(self, 'input', 'output')
+
+  def test_blacklist(self):
+    with utils.tmp_dir():
+      # Create input files and directories.
+      fw = test_utils.FileWriter(os.path.join(os.getcwd(), 'input'))
+      fw.mkdir('.git')
+      fw.write(os.path.join('.git', 'index'))
+      fw.write('somefile')
+      fw.write('.DS_STORE')
+      fw.write('leftover.pyc')
+      fw.write('.pycfile')
+
+      # Zip, unzip.
+      zip_utils.zip('input', 'test.zip', blacklist=['.git', '.DS*', '*.pyc'])
+      zip_utils.unzip('test.zip', 'output')
+
+      # Remove the files/dirs we don't expect to see in output, so that we can
+      # use self._compare_trees to check the results.
+      fw.remove(os.path.join('.git', 'index'))
+      fw.remove('.git')
+      fw.remove('.DS_STORE')
+      fw.remove('leftover.pyc')
+
+      # Compare results.
+      test_utils.compare_trees(self, 'input', 'output')
+
+  def test_nonexistent_dir(self):
+    with utils.tmp_dir():
+      with self.assertRaises(IOError):
+        zip_utils.zip('input', 'test.zip')
+
+
+if __name__ == '__main__':
+  unittest.main()