From 485550989ff0672100aee05005039f045b5f7f2a Mon Sep 17 00:00:00 2001
From: John Sirois <john.sirois@gmail.com>
Date: Mon, 23 Jul 2018 13:19:47 -0600
Subject: [PATCH] Fix pex extraction perms. (#528)

Previously, a zipped pex would lose permission bits when exracted to the
filesystem for `--not-zip-safe` pexes or `PEX_FORCE_LOCAL` runs. This
was due to an underlying bug in the `zipfile` stdlib tracked here:
  https://bugs.python.org/issue15795

Work around the bug in `zipfile.Zipfile` by extending it and running a
chmod'ing cleanup whenever `extract` or `extractall` is called.

Fixes #527
---
 pex/archiver.py      |  4 ++--
 pex/common.py        | 19 +++++++++++++++++-
 tests/test_common.py | 48 ++++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 66 insertions(+), 5 deletions(-)

diff --git a/pex/archiver.py b/pex/archiver.py
index 20baade16..d2afda48e 100644
--- a/pex/archiver.py
+++ b/pex/archiver.py
@@ -6,7 +6,7 @@
 import tarfile
 import zipfile
 
-from .common import safe_mkdtemp
+from .common import PermPreservingZipFile, safe_mkdtemp
 
 
 class Archiver(object):
@@ -19,7 +19,7 @@ class InvalidArchive(Error): pass
     '.tar.gz': (tarfile.TarFile.open, tarfile.ReadError),
     '.tar.bz2': (tarfile.TarFile.open, tarfile.ReadError),
     '.tgz': (tarfile.TarFile.open, tarfile.ReadError),
-    '.zip': (zipfile.ZipFile, zipfile.BadZipfile)
+    '.zip': (PermPreservingZipFile, zipfile.BadZipfile)
   }
 
   @classmethod
diff --git a/pex/common.py b/pex/common.py
index 546e4a31f..f7bb53cf6 100644
--- a/pex/common.py
+++ b/pex/common.py
@@ -78,10 +78,27 @@ def teardown(self):
 _MKDTEMP_SINGLETON = MktempTeardownRegistry()
 
 
+class PermPreservingZipFile(zipfile.ZipFile):
+  """A ZipFile that works around https://bugs.python.org/issue15795"""
+
+  def _extract_member(self, member, targetpath, pwd):
+    result = super(PermPreservingZipFile, self)._extract_member(member, targetpath, pwd)
+    info = member if isinstance(member, zipfile.ZipInfo) else self.getinfo(member)
+    self._chmod(info, result)
+    return result
+
+  def _chmod(self, info, path):
+    # This magic works to extract perm bits from the 32 bit external file attributes field for
+    # unix-created zip files, for the layout, see:
+    #   https://www.forensicswiki.org/wiki/ZIP#External_file_attributes
+    attr = info.external_attr >> 16
+    os.chmod(path, attr)
+
+
 @contextlib.contextmanager
 def open_zip(path, *args, **kwargs):
   """A contextmanager for zip files.  Passes through positional and kwargs to zipfile.ZipFile."""
-  with contextlib.closing(zipfile.ZipFile(path, *args, **kwargs)) as zip:
+  with contextlib.closing(PermPreservingZipFile(path, *args, **kwargs)) as zip:
     yield zip
 
 
diff --git a/tests/test_common.py b/tests/test_common.py
index 643f4b007..23762d6ad 100644
--- a/tests/test_common.py
+++ b/tests/test_common.py
@@ -1,13 +1,14 @@
 # Copyright 2016 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
-
+import contextlib
 import errno
 import os
 from contextlib import contextmanager
 
 import pytest
 
-from pex.common import rename_if_empty
+from pex.common import PermPreservingZipFile, chmod_plus_x, rename_if_empty, touch
+from pex.testing import temporary_dir
 
 try:
   from unittest import mock
@@ -42,3 +43,46 @@ def test_rename_if_empty_enotempty():
 
 def test_rename_if_empty_eperm():
   rename_if_empty_test(errno.EPERM, expect_raises=OSError)
+
+
+def extract_perms(path):
+  return oct(os.stat(path).st_mode)
+
+
+@contextlib.contextmanager
+def zip_fixture():
+  with temporary_dir() as target_dir:
+    one = os.path.join(target_dir, 'one')
+    touch(one)
+
+    two = os.path.join(target_dir, 'two')
+    touch(two)
+    chmod_plus_x(two)
+
+    assert extract_perms(one) != extract_perms(two)
+
+    zip_file = os.path.join(target_dir, 'test.zip')
+    with contextlib.closing(PermPreservingZipFile(zip_file, 'w')) as zf:
+      zf.write(one, 'one')
+      zf.write(two, 'two')
+
+    yield zip_file, os.path.join(target_dir, 'extract'), one, two
+
+
+def test_perm_preserving_zipfile_extractall():
+  with zip_fixture() as (zip_file, extract_dir, one, two):
+    with contextlib.closing(PermPreservingZipFile(zip_file)) as zf:
+      zf.extractall(extract_dir)
+
+      assert extract_perms(one) == extract_perms(os.path.join(extract_dir, 'one'))
+      assert extract_perms(two) == extract_perms(os.path.join(extract_dir, 'two'))
+
+
+def test_perm_preserving_zipfile_extract():
+  with zip_fixture() as (zip_file, extract_dir, one, two):
+    with contextlib.closing(PermPreservingZipFile(zip_file)) as zf:
+      zf.extract('one', path=extract_dir)
+      zf.extract('two', path=extract_dir)
+
+      assert extract_perms(one) == extract_perms(os.path.join(extract_dir, 'one'))
+      assert extract_perms(two) == extract_perms(os.path.join(extract_dir, 'two'))