Fixes #11 Detect and overwrite corrupt PEP 514 registry entries (#21)

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,70 @@ 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)
+                # Continue, hopefully the next new_name is available
+            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 +229,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