diff --git a/_msbuild.py b/_msbuild.py index d1a11c3..270e099 100644 --- a/_msbuild.py +++ b/_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', ) diff --git a/_msbuild_test.py b/_msbuild_test.py index 7053e36..dfe1820 100644 --- a/_msbuild_test.py +++ b/_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', ), ) diff --git a/src/_native/misc.cpp b/src/_native/misc.cpp index f2b8fb6..f3499e7 100644 --- a/src/_native/misc.cpp +++ b/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; +} + } diff --git a/src/manage/pep514utils.py b/src/manage/pep514utils.py index a277544..c89007c 100644 --- a/src/manage/pep514utils.py +++ b/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) diff --git a/tests/test_pep514utils.py b/tests/test_pep514utils.py new file mode 100644 index 0000000..80f8bf7 --- /dev/null +++ b/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