Fixes #11 Detect and overwrite corrupt PEP 514 registry entries

pep514utils now attempts to move existing registry entries if they appear invalid.
This should reduce the likelihood of bad MSI uninstalls causing repeated warnings
for the rest of time, at a small risk of corrupting (intentionally useless) entries.

Author: zooba

_msbuild.py

@@ -83,6 +83,7 @@ class ResourceFile(CSourceFile):
     CFunction('fd_supports_vt100'),
     CFunction('date_as_str'),
     CFunction('datetime_as_str'),
+    CFunction('reg_rename_key'),
     source='src/_native',
     RootNamespace='_native',
 )

_msbuild_test.py

@@ -49,6 +49,7 @@
         CFunction('fd_supports_vt100'),
         CFunction('date_as_str'),
         CFunction('datetime_as_str'),
+        CFunction('reg_rename_key'),
         source='src/_native',
     ),
 )

src/_native/misc.cpp

@@ -95,4 +95,28 @@ PyObject *datetime_as_str(PyObject *, PyObject *, PyObject *) {
     return PyUnicode_FromWideChar(buffer, cch - 1);
 }
 
+PyObject *reg_rename_key(PyObject *, PyObject *args, PyObject *kwargs) {
+    static const char * keywords[] = {"handle", "name1", "name2", NULL};
+    PyObject *handle_obj;
+    wchar_t *name1 = NULL;
+    wchar_t *name2 = NULL;
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO&O&:reg_rename_key", keywords,
+        &handle_obj, as_utf16, &name1, as_utf16, &name2)) {
+        return NULL;
+    }
+    PyObject *r = NULL;
+    HKEY h;
+    if (PyLong_AsNativeBytes(handle_obj, &h, sizeof(h), -1) >= 0) {
+        int err = (int)RegRenameKey(h, name1, name2);
+        if (!err) {
+            r = Py_GetConstant(Py_CONSTANT_NONE);
+        } else {
+            PyErr_SetFromWindowsErr(err);
+        }
+    }
+    PyMem_Free(name1);
+    PyMem_Free(name2);
+    return r;
+}
+
 }

src/manage/pep514utils.py

@@ -136,7 +136,7 @@ def _update_reg_values(key, data, install, exclude=set()):
             winreg.SetValueEx(key, k, None, v_kind, v)
 
 
-def _is_tag_managed(company_key, tag_name):
+def _is_tag_managed(company_key, tag_name, *, creating=False):
     try:
         tag = winreg.OpenKey(company_key, tag_name)
     except FileNotFoundError:
@@ -147,7 +147,69 @@ def _is_tag_managed(company_key, tag_name):
                 return True
         except FileNotFoundError:
             pass
-    return False
+        if not creating:
+            return False
+
+        # gh-11: Clean up invalid entries from other installers
+        # It's highly likely that our old MSI installer wouldn't properly remove
+        # its registry keys on uninstall, so we'll check for the InstallPath
+        # subkey and if it's missing, back up the key and then use it ourselves.
+        try:
+            with _reg_open(tag, "InstallPath") as subkey:
+                if subkey:
+                    # if InstallPath refers to a directory that exists,
+                    # leave it alone.
+                    p = winreg.QueryValueEx(subkey, None)[0]
+                    if p and Path(p).exists():
+                        return False
+        except FileNotFoundError:
+            pass
+        except OSError:
+            # If we couldn't access it for some reason, leave it alone.
+            return False
+
+    # Existing key is almost certainly not valid, so let's rename it,
+    # warn the user, and then continue.
+    LOGGER.debug("Key %s appears invalid, so moving it and taking it for this "
+                 "new install", tag_name)
+    try:
+        from _native import reg_rename_key
+    except ImportError:
+        LOGGER.debug("Failed to import reg_rename_key", exc_info=True)
+        return False
+
+    parent_name, _, orig_name = tag_name.replace("/", "\\").rpartition("\\")
+    with _reg_open(company_key, parent_name, writable=True) as tag:
+        if not tag:
+            # Key is no longer there, so we can use it
+            return True
+        for i in range(1000):
+            try:
+                new_name = f"{orig_name}.{i}"
+                # Raises standard PermissionError (5) if new_name exists
+                reg_rename_key(tag.handle, orig_name, new_name)
+                LOGGER.warn("An existing registry key for %s was renamed to %s "
+                            "because it appeared to be invalid. If this is "
+                            "correct, the registry key can be safely deleted. "
+                            "To avoid this in future, ensure that the "
+                            "InstallPath key refers to a valid path.",
+                            tag_name, new_name)
+                break
+            except FileNotFoundError:
+                LOGGER.debug("Original key disappeared, so we will claim it")
+                return True
+            except PermissionError:
+                LOGGER.debug("Failed to rename %s to %s", orig_name, new_name,
+                             exc_info=True)
+            except OSError:
+                LOGGER.debug("Unexpected error while renaming %s to %s",
+                             orig_name, new_name, exc_info=True)
+                raise
+        else:
+            LOGGER.warn("Attempted to clean up invalid registry key %s but "
+                        "failed after too many attempts.", tag_name);
+            return False
+        return True
 
 
 def _split_root(root_name):
@@ -166,7 +228,7 @@ def _split_root(root_name):
 def update_registry(root_name, install, data):
     hive, name = _split_root(root_name)
     with winreg.CreateKey(hive, name) as root:
-        if _is_tag_managed(root, data["Key"]):
+        if _is_tag_managed(root, data["Key"], creating=True):
             with winreg.CreateKey(root, data["Key"]) as tag:
                 LOGGER.debug("Creating/updating %s\\%s", root_name, data["Key"])
                 winreg.SetValueEx(tag, "ManagedByPyManager", None, winreg.REG_DWORD, 1)

tests/test_pep514utils.py

@@ -0,0 +1,54 @@
+import pytest
+import winreg
+
+from manage import pep514utils
+
+REG_TEST_ROOT = r"Software\Python\PyManagerTesting"
+
+@pytest.fixture(scope='function')
+def registry():
+    try:
+        with winreg.CreateKey(winreg.HKEY_CURRENT_USER, REG_TEST_ROOT) as key:
+            yield key
+    finally:
+        pep514utils._reg_rmtree(winreg.HKEY_CURRENT_USER, REG_TEST_ROOT)
+
+
+def init_reg(registry, **keys):
+    for k, v in keys.items():
+        if isinstance(v, dict):
+            with winreg.CreateKey(registry, k) as subkey:
+                init_reg(subkey, **v)
+        elif isinstance(v, str):
+            winreg.SetValueEx(registry, k, None, winreg.REG_SZ, v)
+        elif isinstance(v, (bytes, bytearray)):
+            winreg.SetValueEx(registry, k, None, winreg.REG_BINARY, v)
+        elif isinstance(v, int):
+            if v.bit_count() < 32:
+                winreg.SetValueEx(registry, k, None, winreg.REG_DWORD, v)
+            else:
+                winreg.SetValueEx(registry, k, None, winreg.REG_QWORD, v)
+        else:
+            raise TypeError("unsupported type in registry")
+
+
+def test_is_tag_managed(registry, tmp_path):
+    init_reg(registry, Company={
+        "1.0": {"InstallPath": {"": str(tmp_path)}},
+        "2.0": {"InstallPath": {"": str(tmp_path)}, "ManagedByPyManager": 0},
+        "2.1": {"InstallPath": {"": str(tmp_path)}, "ManagedByPyManager": 1},
+        "3.0": {"InstallPath": {"": str(tmp_path / "missing")}},
+        "3.0.0": {"": "Just in the way here"},
+        "3.0.1": {"": "Also in the way here"},
+    })
+
+    assert not pep514utils._is_tag_managed(registry, r"Company\1.0")
+    assert not pep514utils._is_tag_managed(registry, r"Company\2.0")
+    assert pep514utils._is_tag_managed(registry, r"Company\2.1")
+
+    assert not pep514utils._is_tag_managed(registry, r"Company\3.0")
+    with pytest.raises(FileNotFoundError):
+        winreg.OpenKey(registry, r"Company\3.0.2")
+    assert pep514utils._is_tag_managed(registry, r"Company\3.0", creating=True)
+    with winreg.OpenKey(registry, r"Company\3.0.2"):
+        pass