gh-300: Improve purge handling of unrecognized files

Author: Punisheroot

src/manage/fsutils.py

@@ -21,13 +21,16 @@ def ensure_tree(path, overwrite_files=True):
         path.parent.mkdir(parents=True, exist_ok=True)
 
 
-def _rglob(root):
+def _rglob(root, follow_links=True):
     q = [root]
     while q:
         r = q.pop(0)
         for f in os.scandir(r):
             p = r / f.name
-            if f.is_dir():
+            is_dir = f.is_dir(follow_symlinks=follow_links)
+            if not follow_links and f.is_junction():
+                is_dir = False
+            if is_dir:
                 q.append(p)
                 yield p, None
             else:
@@ -92,6 +95,13 @@ def rmtree(path, after_5s_warning=None, remove_ext_first=()):
 
     if isinstance(path, (str, bytes)):
         path = Path(path)
+    if os.path.islink(path):
+        unlink(path, after_5s_warning=after_5s_warning)
+        return
+    if os.path.isjunction(path):
+        LOGGER.debug("Removing junction without traversing it: %s", path)
+        _rmdir(path, on_fail=lambda p: LOGGER.warn("Failed to remove %s", p))
+        return
     if not path.is_dir():
         if path.is_file():
             unlink(path)
@@ -131,7 +141,7 @@ def rmtree(path, after_5s_warning=None, remove_ext_first=()):
 
     to_rmdir = [path]
     to_unlink = []
-    for d, f in _rglob(path):
+    for d, f in _rglob(path, follow_links=False):
         if after_5s_warning and (time.monotonic() - start) > 5:
             LOGGER.warn(after_5s_warning)
             after_5s_warning = None

src/manage/uninstall_command.py

@@ -86,6 +86,7 @@ def execute(cmd):
         if not cmd.ask_yn("Uninstall all runtimes?"):
             LOGGER.debug("END uninstall_command.execute")
             return
+        installs_in_use = set()
         for i in installed:
             LOGGER.info("Purging %s from %s", i["display-name"], i["prefix"])
             try:
@@ -97,6 +98,7 @@ def execute(cmd):
             except FilesInUseError:
                 LOGGER.warn("Unable to purge %s because it is still in use.",
                             i["display-name"])
+                installs_in_use.add(Path(i["prefix"]))
                 continue
         LOGGER.info("Purging saved downloads from %s", cmd.download_dir)
         rmtree(cmd.download_dir, after_5s_warning=warn_msg.format("cached downloads"))
@@ -106,6 +108,14 @@ def execute(cmd):
         for _, cleanup in SHORTCUT_HANDLERS.values():
             if cleanup:
                 cleanup(cmd, [])
+        unknown = [p for p in _iterdir(cmd.install_dir)
+                   if p not in installs_in_use]
+        if unknown:
+            LOGGER.info("Purging unrecognized files from %s", cmd.install_dir)
+            for p in unknown:
+                rmtree(p, after_5s_warning=warn_msg.format(
+                    "unrecognized files"
+                ))
         LOGGER.debug("END uninstall_command.execute")
         return
 

tests/test_fsutils.py

@@ -1,5 +1,6 @@
 import pytest
 import shutil
+import _winapi
 
 from copy import copy
 
@@ -58,6 +59,36 @@ def test_rmtree(tree):
     assert not tree.exists()
 
 
+def test_rmtree_junction(tmp_path):
+    target = tmp_path / "target"
+    target.mkdir()
+    target_file = target / "preserve.txt"
+    target_file.write_bytes(b"preserve")
+    junction = tmp_path / "junction"
+    _winapi.CreateJunction(str(target), str(junction))
+
+    rmtree(junction)
+
+    assert not junction.exists()
+    assert target_file.read_bytes() == b"preserve"
+
+
+def test_rmtree_nested_junction(tmp_path):
+    target = tmp_path / "target"
+    target.mkdir()
+    target_file = target / "preserve.txt"
+    target_file.write_bytes(b"preserve")
+    root = tmp_path / "root"
+    root.mkdir()
+    junction = root / "junction"
+    _winapi.CreateJunction(str(target), str(junction))
+
+    rmtree(root)
+
+    assert not root.exists()
+    assert target_file.read_bytes() == b"preserve"
+
+
 def test_atomic_unlink(tree):
     files = [tree / "c/d", tree / "b"]
     assert all([f.is_file() for f in files])

tests/test_uninstall_command.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 
 from manage import uninstall_command as UC
+from manage.exceptions import FilesInUseError
 
 
 def test_purge_global_dir(monkeypatch, registry, tmp_path):
@@ -25,3 +26,55 @@ def test_null_purge(fake_config):
     cmd.confirm = False
     cmd.purge = True
     UC.execute(cmd)
+
+
+def test_purge_unknown_files(fake_config):
+    cmd = fake_config
+    cmd.args = ["--purge"]
+    cmd.confirm = False
+    cmd.purge = True
+
+    unknown_file = cmd.install_dir / "unknown.txt"
+    unknown_file.write_bytes(b"unknown")
+    broken_runtime = cmd.install_dir / "broken-runtime"
+    broken_runtime.mkdir()
+    (broken_runtime / "__install__.json").write_text("invalid")
+
+    UC.execute(cmd)
+
+    assert not unknown_file.exists()
+    assert not broken_runtime.exists()
+
+
+def test_purge_preserves_runtime_in_use(fake_config, monkeypatch):
+    cmd = fake_config
+    cmd.args = ["--purge"]
+    cmd.confirm = False
+    cmd.purge = True
+
+    runtime = cmd.install_dir / "runtime"
+    runtime.mkdir()
+    executable = runtime / "python.exe"
+    executable.write_bytes(b"in use")
+    cmd.installs = [{
+        "display-name": "Runtime in use",
+        "prefix": runtime,
+    }]
+
+    unknown_dir = cmd.install_dir / "unknown"
+    unknown_dir.mkdir()
+    (unknown_dir / "file.txt").write_bytes(b"unknown")
+
+    rmtree = UC.rmtree
+
+    def locked_rmtree(path, *args, **kwargs):
+        if Path(path) == runtime:
+            raise FilesInUseError([executable])
+        return rmtree(path, *args, **kwargs)
+
+    monkeypatch.setattr(UC, "rmtree", locked_rmtree)
+
+    UC.execute(cmd)
+
+    assert executable.is_file()
+    assert not unknown_dir.exists()