#4489: Fix usage of fd-based functions to new api introduced earlier today

Also add an explicit test for safe implementation usage on supported platforms.

As a side effect, this commit adds a module-level attribute 'rmtree_is_safe'
which offers introspection whether the current rmtree implementation is safe
against symlink attacks.

3 files changed, 44 insertions, 17 deletions

diff --git a/Doc/library/shutil.rst b/Doc/library/shutil.rst
index c3eb990..5521841 100644
--- a/Doc/library/shutil.rst
+++ b/Doc/library/shutil.rst
@@ -195,9 +195,9 @@ Directory and files operations
       The default :func:`rmtree` function is susceptible to a symlink attack:
       given proper timing and circumstances, attackers can use it to delete
       files they wouldn't be able to access otherwise.  Thus -- on platforms
-      that support the necessary fd-based functions :func:`os.openat` and
-      :func:`os.unlinkat` -- a safe version of :func:`rmtree` is used, which
-      isn't vulnerable.
+      that support the necessary fd-based functions -- a safe version of
+      :func:`rmtree` is used, which isn't vulnerable. In this case
+      :data:`rmtree_is_safe` is set to True.
 
    If *onerror* is provided, it must be a callable that accepts three
    parameters: *function*, *path*, and *excinfo*.
@@ -210,8 +210,15 @@ Directory and files operations
 
    .. versionchanged:: 3.3
       Added a safe version that is used automatically if platform supports
-      the fd-based functions :func:`os.openat` and :func:`os.unlinkat`.
+      fd-based functions.
 
+.. data:: rmtree_is_safe
+
+   Indicates whether the current platform and implementation has a symlink
+   attack-proof version of :func:`rmtree`. Currently this is only true for
+   platforms supporting fd-based directory access functions.
+
+   .. versionadded:: 3.3
 
 .. function:: move(src, dst)
 
diff --git a/Lib/shutil.py b/Lib/shutil.py
index 1b05484..c41046a 100644
--- a/Lib/shutil.py
+++ b/Lib/shutil.py
@@ -362,9 +362,9 @@ def _rmtree_unsafe(path, onerror):
             _rmtree_unsafe(fullname, onerror)
         else:
             try:
-                os.remove(fullname)
+                os.unlink(fullname)
             except os.error:
-                onerror(os.remove, fullname, sys.exc_info())
+                onerror(os.unlink, fullname, sys.exc_info())
     try:
         os.rmdir(path)
     except os.error:
@@ -374,21 +374,21 @@ def _rmtree_unsafe(path, onerror):
 def _rmtree_safe_fd(topfd, path, onerror):
     names = []
     try:
-        names = os.flistdir(topfd)
+        names = os.listdir(topfd)
     except os.error:
-        onerror(os.flistdir, path, sys.exc_info())
+        onerror(os.listdir, path, sys.exc_info())
     for name in names:
         fullname = os.path.join(path, name)
         try:
-            orig_st = os.fstatat(topfd, name)
+            orig_st = os.stat(name, dir_fd=topfd)
             mode = orig_st.st_mode
         except os.error:
             mode = 0
         if stat.S_ISDIR(mode):
             try:
-                dirfd = os.openat(topfd, name, os.O_RDONLY)
+                dirfd = os.open(name, os.O_RDONLY, dir_fd=topfd)
             except os.error:
-                onerror(os.openat, fullname, sys.exc_info())
+                onerror(os.open, fullname, sys.exc_info())
             else:
                 try:
                     if os.path.samestat(orig_st, os.fstat(dirfd)):
@@ -397,21 +397,22 @@ def _rmtree_safe_fd(topfd, path, onerror):
                     os.close(dirfd)
         else:
             try:
-                os.unlinkat(topfd, name)
+                os.unlink(name, dir_fd=topfd)
             except os.error:
-                onerror(os.unlinkat, fullname, sys.exc_info())
+                onerror(os.unlink, fullname, sys.exc_info())
     try:
         os.rmdir(path)
     except os.error:
         onerror(os.rmdir, path, sys.exc_info())
 
-_use_fd_functions = hasattr(os, 'openat') and hasattr(os, 'unlinkat')
+rmtree_is_safe = _use_fd_functions = (os.unlink in os.supports_dir_fd and
+                                      os.open in os.supports_dir_fd)
 def rmtree(path, ignore_errors=False, onerror=None):
     """Recursively delete a directory tree.
 
     If ignore_errors is set, errors are ignored; otherwise, if onerror
     is set, it is called to handle the error with arguments (func,
-    path, exc_info) where func is os.listdir, os.remove, or os.rmdir;
+    path, exc_info) where func is platform and implementation dependent;
     path is the argument to that function that caused it to fail; and
     exc_info is a tuple returned by sys.exc_info().  If ignore_errors
     is false and onerror is None, an exception is raised.
diff --git a/Lib/test/test_shutil.py b/Lib/test/test_shutil.py
index 9c0c52c..b12d259 100644
--- a/Lib/test/test_shutil.py
+++ b/Lib/test/test_shutil.py
@@ -159,8 +159,7 @@ class TestShutil(unittest.TestCase):
         # at os.listdir.  The first failure may legally
         # be either.
         if 0 <= self.errorState < 2:
-            if (func is os.remove or
-                hasattr(os, 'unlinkat') and func is os.unlinkat):
+            if func is os.unlink:
                 self.assertIn(arg, [self.child_file_path, self.child_dir_path])
             else:
                 if self.errorState == 1:
@@ -486,6 +485,26 @@ class TestShutil(unittest.TestCase):
         shutil.copyfile(link, dst)
         self.assertFalse(os.path.islink(dst))
 
+    def test_rmtree_uses_safe_fd_version_if_available(self):
+        if os.unlink in os.supports_dir_fd and os.open in os.supports_dir_fd:
+            self.assertTrue(shutil._use_fd_functions)
+            self.assertTrue(shutil.rmtree_is_safe)
+            tmp_dir = self.mkdtemp()
+            d = os.path.join(tmp_dir, 'a')
+            os.mkdir(d)
+            try:
+                real_rmtree = shutil._rmtree_safe_fd
+                class Called(Exception): pass
+                def _raiser(*args, **kwargs):
+                    raise Called
+                shutil._rmtree_safe_fd = _raiser
+                self.assertRaises(Called, shutil.rmtree, d)
+            finally:
+                shutil._rmtree_safe_fd = real_rmtree
+        else:
+            self.assertFalse(shutil._use_fd_functions)
+            self.assertFalse(shutil.rmtree_is_safe)
+
     def test_rmtree_dont_delete_file(self):
         # When called on a file instead of a directory, don't delete it.
         handle, path = tempfile.mkstemp()