Update sys.getwindowsversion to use better calculation

python · zooba · Sep 7, 2022 · Aug 25, 2022 · Aug 30, 2022 · Aug 30, 2022
commit aef3a4d78710bade3bf080fce905d6e2dd38ff39
diff --git a/Lib/ntpath.py b/Lib/ntpath.py
@@ -732,9 +732,8 @@ def realpath(path, *, strict=False):
         return path
 
 
-# Win9x family and earlier have no Unicode filename support.
-supports_unicode_filenames = (hasattr(sys, "getwindowsversion") and
-                              sys.getwindowsversion()[3] >= 2)
+# All supported version have Unicode filename support.
+supports_unicode_filenames = True
 
 def relpath(path, start=None):
     """Return a relative version of a path"""

@@ -1486,6 +1486,108 @@ static PyStructSequence_Desc windows_version_desc = {
                                       via indexing, the rest are name only */
 };
 
+static PyObject *
+_sys_getwindowsversion_from_wmi()
+{
+    PyObject *wmi_module = PyImport_ImportModule("_wmi");
+    if (!wmi_module) {
+        return NULL;
+    }
+
+    PyObject *queryResult = PyObject_CallMethod(wmi_module, "exec_query", "s",
+        "SELECT Version FROM Win32_OperatingSystem"
+    );
+    Py_DECREF(wmi_module);
+
+    if (!queryResult) {
+        return NULL;
+    }
+
+    char versionString[256];
+    const char *queryResultString = PyUnicode_AsUTF8(queryResult);
+    if (!queryResultString) {
+        return NULL;
+    }
+    if (strcpy_s(versionString, sizeof(versionString), queryResultString)) {
+        PyErr_SetFromErrno(0);
+        Py_DECREF(queryResultString);
+        return NULL;
+    }
+    if (0 != _strnicmp(versionString, "version=", 8)) {
+        PyErr_Format(PyExc_SystemError, "invalid WMI result: %s", versionString);
+        Py_DECREF(queryResult);
+        return NULL;
+    }
+    char *majorStart = strchr(versionString, '=');
+    if (!majorStart || !majorStart[1]) {
+        PyErr_Format(PyExc_SystemError, "invalid WMI result: %s", versionString);
+        Py_DECREF(queryResult);
+        return NULL;
+    }
+    majorStart += 1;
+    char *minorStart = strchr(majorStart, '.');
+    if (!minorStart || !minorStart[1]) {
+        PyErr_Format(PyExc_SystemError, "invalid WMI result: %s", versionString);
+        Py_DECREF(queryResult);
+        return NULL;
+    }
+    *minorStart = '\0';
+    minorStart += 1;
+    char *buildStart = strchr(minorStart, '.');
+    if (!buildStart || !buildStart[1]) {
+        PyErr_Format(PyExc_SystemError, "invalid WMI result: %s", versionString);
+        Py_DECREF(queryResult);
+        return NULL;
+    }
+    *buildStart = '\0';
+    buildStart += 1;
+    PyObject *result = Py_BuildValue("(NNN)",
+        PyLong_FromString(majorStart, NULL, 10),
+        PyLong_FromString(minorStart, NULL, 10),
+        PyLong_FromString(buildStart, NULL, 10)
+    );
+    Py_DECREF(queryResult);
+    return result;
+}
+
+static PyObject *
+_sys_getwindowsversion_from_kernel32()
+{
+    HANDLE hKernel32;
+    wchar_t kernel32_path[MAX_PATH];
+    LPVOID verblock;
+    DWORD verblock_size;
+    VS_FIXEDFILEINFO *ffi;
+    UINT ffi_len;
+    DWORD realMajor, realMinor, realBuild;
+
+    Py_BEGIN_ALLOW_THREADS
+    hKernel32 = GetModuleHandleW(L"kernel32.dll");
+    Py_END_ALLOW_THREADS
+    if (!hKernel32 || !GetModuleFileNameW(hKernel32, kernel32_path, MAX_PATH)) {
+        PyErr_SetFromWindowsErr(0);
+        return NULL;
+    }
+    verblock_size = GetFileVersionInfoSizeW(kernel32_path, NULL);
+    if (!verblock_size) {
+        PyErr_SetFromWindowsErr(0);
+        return NULL;
+    }
+    verblock = PyMem_RawMalloc(verblock_size);
+    if (!verblock ||
+        !GetFileVersionInfoW(kernel32_path, 0, verblock_size, verblock) ||
+        !VerQueryValueW(verblock, L"", (LPVOID)&ffi, &ffi_len)) {
+        PyErr_SetFromWindowsErr(0);
+        return NULL;
+    }
+
+    realMajor = HIWORD(ffi->dwProductVersionMS);
+    realMinor = LOWORD(ffi->dwProductVersionMS);
+    realBuild = HIWORD(ffi->dwProductVersionLS);
+    PyMem_RawFree(verblock);
+    return Py_BuildValue("(kkk)", realMajor, realMinor, realBuild);
+}
+
 /* Disable deprecation warnings about GetVersionEx as the result is
    being passed straight through to the caller, who is responsible for
    using it correctly. */
@@ -1515,11 +1617,13 @@ sys_getwindowsversion_impl(PyObject *module)
     PyObject *version;
     int pos = 0;
     OSVERSIONINFOEXW ver;
-    DWORD realMajor, realMinor, realBuild;
-    HANDLE hKernel32;
-    wchar_t kernel32_path[MAX_PATH];
-    LPVOID verblock;
-    DWORD verblock_size;
+
+    version = PyObject_GetAttrString(module, "_cached_windows_version");
+    if (version && PyObject_TypeCheck(version, &WindowsVersionType)) {
+        return version;
+    }
+    Py_XDECREF(version);
+    PyErr_Clear();
 
     ver.dwOSVersionInfoSize = sizeof(ver);
     if (!GetVersionExW((OSVERSIONINFOW*) &ver))
@@ -1539,41 +1643,39 @@ sys_getwindowsversion_impl(PyObject *module)
     PyStructSequence_SET_ITEM(version, pos++, PyLong_FromLong(ver.wSuiteMask));
     PyStructSequence_SET_ITEM(version, pos++, PyLong_FromLong(ver.wProductType));
 
-    realMajor = ver.dwMajorVersion;
-    realMinor = ver.dwMinorVersion;
-    realBuild = ver.dwBuildNumber;
-
     // GetVersion will lie if we are running in a compatibility mode.
     // We need to read the version info from a system file resource
     // to accurately identify the OS version. If we fail for any reason,
     // just return whatever GetVersion said.
-    Py_BEGIN_ALLOW_THREADS
-    hKernel32 = GetModuleHandleW(L"kernel32.dll");
-    Py_END_ALLOW_THREADS
-    if (hKernel32 && GetModuleFileNameW(hKernel32, kernel32_path, MAX_PATH) &&
-        (verblock_size = GetFileVersionInfoSizeW(kernel32_path, NULL)) &&
-        (verblock = PyMem_RawMalloc(verblock_size))) {
-        VS_FIXEDFILEINFO *ffi;
-        UINT ffi_len;
-
-        if (GetFileVersionInfoW(kernel32_path, 0, verblock_size, verblock) &&
-            VerQueryValueW(verblock, L"", (LPVOID)&ffi, &ffi_len)) {
-            realMajor = HIWORD(ffi->dwProductVersionMS);
-            realMinor = LOWORD(ffi->dwProductVersionMS);
-            realBuild = HIWORD(ffi->dwProductVersionLS);
+    PyObject *realVersion = _sys_getwindowsversion_from_wmi();
+    if (!realVersion) {
+        return NULL;
+        PyErr_Clear();
+        realVersion = _sys_getwindowsversion_from_kernel32();
+        if (!realVersion) {
+            PyErr_Clear();
+            realVersion = Py_BuildValue("(kkk)",
+                ver.dwMajorVersion,
+                ver.dwMinorVersion,
+                ver.dwBuildNumber
+            );
         }
-        PyMem_RawFree(verblock);
     }
-    PyStructSequence_SET_ITEM(version, pos++, Py_BuildValue("(kkk)",
-        realMajor,
-        realMinor,
-        realBuild
-    ));
+
+    if (realVersion) {
+        PyStructSequence_SET_ITEM(version, pos++, realVersion);
+    }
 
     if (PyErr_Occurred()) {
         Py_DECREF(version);
         return NULL;
     }
+
+    if (PyObject_SetAttrString(module, "_cached_windows_version", version) < 0) {
+        Py_DECREF(version);
+        return NULL;
+    }
+
     return version;
 }