Merge branch 'master' of github.com:giampaolo/psutil

9 files changed, 135 insertions, 82 deletions

diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
index 3908884e..c0d73645 100644
--- a/psutil/_psutil_windows.c
+++ b/psutil/_psutil_windows.c
@@ -3758,7 +3758,7 @@ void init_psutil_windows(void)
     // set SeDebug for the current process
     psutil_set_se_debug();
     psutil_setup();
-    if (psutil_loadlibs() != 0)
+    if (psutil_load_globals() != 0)
         return NULL;
 
 #if PY_MAJOR_VERSION >= 3
diff --git a/psutil/arch/windows/global.c b/psutil/arch/windows/global.c
index bb9970fe..da3bdc6e 100644
--- a/psutil/arch/windows/global.c
+++ b/psutil/arch/windows/global.c
@@ -2,6 +2,10 @@
  * Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
  * Use of this source code is governed by a BSD-style license that can be
  * found in the LICENSE file.
+ *
+ * This code is executed on import. It loads private/undocumented
+ * Windows APIs and sets Windows version constants so that they are
+ * available globally.
  */
 
 #include <windows.h>
@@ -10,6 +14,9 @@
 #include "global.h"
 
 
+// Needed to make this globally visible.
+int PSUTIL_WINVER;
+
 // A wrapper around GetModuleHandle and GetProcAddress.
 PVOID
 psutil_GetProcAddress(LPCSTR libname, LPCSTR procname) {
@@ -52,15 +59,11 @@ psutil_GetProcAddressFromLib(LPCSTR libname, LPCSTR procname) {
 }
 
 
-/*
- * This is executed on import and loads Windows APIs so that they
- * are available globally.
- */
+static int
 psutil_loadlibs() {
     /*
      * Mandatory.
      */
-
     psutil_NtQuerySystemInformation = psutil_GetProcAddressFromLib(
         "ntdll.dll", "NtQuerySystemInformation");
     if (psutil_NtQuerySystemInformation == NULL)
@@ -108,10 +111,14 @@ psutil_loadlibs() {
     if (! psutil_GetExtendedUdpTable)
         return 1;
 
+    psutil_RtlGetVersion = psutil_GetProcAddressFromLib(
+        "ntdll.dll", "RtlGetVersion");
+    if (! psutil_RtlGetVersion)
+        return 1;
+
     /*
      * Optional.
      */
-
     // minimum requirement: Win Vista
     psutil_GetTickCount64 = psutil_GetProcAddress(
         "kernel32", "GetTickCount64");
@@ -127,3 +134,44 @@ psutil_loadlibs() {
     PyErr_Clear();
     return 0;
 }
+
+
+static int
+psutil_set_winver() {
+    RTL_OSVERSIONINFOEXW versionInfo;
+    ULONG maj;
+    ULONG min;
+
+    versionInfo.dwOSVersionInfoSize = sizeof(RTL_OSVERSIONINFOEXW);
+    memset(&versionInfo, 0, sizeof(RTL_OSVERSIONINFOEXW));
+    psutil_RtlGetVersion((PRTL_OSVERSIONINFOW)&versionInfo);
+    maj = versionInfo.dwMajorVersion;
+    min = versionInfo.dwMinorVersion;
+    if (maj == 5 && min == 1)
+        PSUTIL_WINVER = PSUTIL_WINDOWS_XP;
+    else if (maj == 5 && min == 2)
+        PSUTIL_WINVER = PSUTIL_WINDOWS_SERVER_2003;
+    else if (maj == 6 && min == 0)
+        PSUTIL_WINVER = PSUTIL_WINDOWS_VISTA;  // or Server 2008
+    else if (maj == 6 && min == 1)
+        PSUTIL_WINVER = PSUTIL_WINDOWS_7;
+    else if (maj == 6 && min == 2)
+        PSUTIL_WINVER = PSUTIL_WINDOWS_8;
+    else if (maj == 6 && min == 3)
+        PSUTIL_WINVER = PSUTIL_WINDOWS_8_1;
+    else if (maj == 10 && min == 0)
+        PSUTIL_WINVER = PSUTIL_WINDOWS_10;
+    else
+        PSUTIL_WINVER = PSUTIL_WINDOWS_NEW;
+    return 0;
+}
+
+
+int
+psutil_load_globals() {
+    if (psutil_loadlibs() != 0)
+        return 1;
+    if (psutil_set_winver() != 0)
+        return 1;
+    return 0;
+}
diff --git a/psutil/arch/windows/global.h b/psutil/arch/windows/global.h
index 8d828776..7bf8b5ca 100644
--- a/psutil/arch/windows/global.h
+++ b/psutil/arch/windows/global.h
@@ -2,10 +2,62 @@
  * Copyright (c) 2009, Giampaolo Rodola'. All rights reserved.
  * Use of this source code is governed by a BSD-style license that can be
  * found in the LICENSE file.
+
+ * List of constants and objects that are globally available.
  */
 
 #include <windows.h>
+#include "ntextapi.h"
+
+extern int PSUTIL_WINVER;
+#define PSUTIL_WINDOWS_XP 51
+#define PSUTIL_WINDOWS_SERVER_2003 52
+#define PSUTIL_WINDOWS_VISTA 60
+#define PSUTIL_WINDOWS_7 61
+#define PSUTIL_WINDOWS_8 62
+#define PSUTIL_WINDOWS_8_1 63
+#define PSUTIL_WINDOWS_10 100
+#define PSUTIL_WINDOWS_NEW MAXLONG
 
-int psutil_loadlibs();
+int psutil_load_globals();
 PVOID psutil_GetProcAddress(LPCSTR libname, LPCSTR procname);
 PVOID psutil_GetProcAddressFromLib(LPCSTR libname, LPCSTR procname);
+
+_NtQuerySystemInformation \
+    psutil_NtQuerySystemInformation;
+
+_NtQueryInformationProcess \
+    psutil_NtQueryInformationProcess;
+
+_NtSetInformationProcess
+    psutil_NtSetInformationProcess;
+
+_WinStationQueryInformationW \
+    psutil_WinStationQueryInformationW;
+
+_RtlIpv4AddressToStringA \
+    psutil_rtlIpv4AddressToStringA;
+
+_RtlIpv6AddressToStringA \
+    psutil_rtlIpv6AddressToStringA;
+
+_GetExtendedTcpTable \
+    psutil_GetExtendedTcpTable;
+
+_GetExtendedUdpTable \
+    psutil_GetExtendedUdpTable;
+
+_GetActiveProcessorCount \
+    psutil_GetActiveProcessorCount;
+
+_GetTickCount64 \
+    psutil_GetTickCount64;
+
+_NtQueryObject \
+    psutil_NtQueryObject;
+
+_GetLogicalProcessorInformationEx \
+    psutil_GetLogicalProcessorInformationEx;
+
+_RtlGetVersion \
+    psutil_RtlGetVersion;
diff --git a/psutil/arch/windows/inet_ntop.c b/psutil/arch/windows/inet_ntop.c
index 4b6c1dfe..3db507b9 100644
--- a/psutil/arch/windows/inet_ntop.c
+++ b/psutil/arch/windows/inet_ntop.c
@@ -9,10 +9,7 @@
 
 // From: https://memset.wordpress.com/2010/10/09/inet_ntop-for-win32/
 PCSTR WSAAPI
-inet_ntop(__in INT family,
-          __in PVOID pAddr,
-          __out_ecount(StringBufSize) PSTR pStringBuf,
-          __in size_t StringBufSize) {
+inet_ntop(INT family, PVOID pAddr, PSTR stringBuf, size_t strBufSize) {
     DWORD dwAddressLength = 0;
     struct sockaddr_storage srcaddr;
     struct sockaddr_in *srcaddr4 = (struct sockaddr_in*) &srcaddr;
@@ -34,13 +31,15 @@ inet_ntop(__in INT family,
         return NULL;
     }
 
-    if (WSAAddressToString((LPSOCKADDR) &srcaddr,
-                           dwAddressLength,
-                           0,
-                           pStringBuf,
-                           (LPDWORD) &StringBufSize) != 0) {
+    if (WSAAddressToStringA(
+            (LPSOCKADDR) &srcaddr,
+            dwAddressLength,
+            0,
+            stringBuf,
+            (LPDWORD) &strBufSize) != 0)
+    {
         PyErr_SetExcFromWindowsErr(PyExc_OSError, WSAGetLastError());
         return NULL;
     }
-    return pStringBuf;
+    return stringBuf;
 }
diff --git a/psutil/arch/windows/inet_ntop.h b/psutil/arch/windows/inet_ntop.h
index 0dbf28bf..2d86c26c 100644
--- a/psutil/arch/windows/inet_ntop.h
+++ b/psutil/arch/windows/inet_ntop.h
@@ -4,6 +4,8 @@
  * found in the LICENSE file.
  */
 
+// because of WSAAddressToStringA
+#define _WINSOCK_DEPRECATED_NO_WARNINGS
 #include <ws2tcpip.h>
 
 PCSTR WSAAPI
diff --git a/psutil/arch/windows/ntextapi.h b/psutil/arch/windows/ntextapi.h
index e0006c7b..d6030caa 100644
--- a/psutil/arch/windows/ntextapi.h
+++ b/psutil/arch/windows/ntextapi.h
@@ -400,22 +400,19 @@ typedef NTSTATUS (NTAPI *_NtQueryInformationProcess)(
     DWORD ProcessInformationClass,
     PVOID ProcessInformation,
     DWORD ProcessInformationLength,
-    PDWORD ReturnLength
-);
+    PDWORD ReturnLength);
 
 typedef NTSTATUS (NTAPI *_NtQuerySystemInformation)(
     ULONG SystemInformationClass,
     PVOID SystemInformation,
     ULONG SystemInformationLength,
-    PULONG ReturnLength
-);
+    PULONG ReturnLength);
 
 typedef NTSTATUS (NTAPI *_NtSetInformationProcess)(
     HANDLE ProcessHandle,
     DWORD ProcessInformationClass,
     PVOID ProcessInformation,
-    DWORD ProcessInformationLength
-);
+    DWORD ProcessInformationLength);
 
 typedef PSTR (NTAPI * _RtlIpv4AddressToStringA)(
     struct in_addr *Addr,
@@ -431,8 +428,7 @@ typedef DWORD (WINAPI * _GetExtendedTcpTable)(
     BOOL bOrder,
     ULONG ulAf,
     TCP_TABLE_CLASS TableClass,
-    ULONG Reserved
-);
+    ULONG Reserved);
 
 typedef DWORD (WINAPI * _GetExtendedUdpTable)(
     PVOID pUdpTable,
@@ -440,8 +436,7 @@ typedef DWORD (WINAPI * _GetExtendedUdpTable)(
     BOOL bOrder,
     ULONG ulAf,
     UDP_TABLE_CLASS TableClass,
-    ULONG Reserved
-);
+    ULONG Reserved);
 
 typedef DWORD (CALLBACK *_GetActiveProcessorCount)(
     WORD GroupNumber);
@@ -454,49 +449,10 @@ typedef NTSTATUS (NTAPI *_NtQueryObject)(
     OBJECT_INFORMATION_CLASS ObjectInformationClass,
     PVOID ObjectInformation,
     ULONG ObjectInformationLength,
-    PULONG ReturnLength
-);
-
-/*
- * ================================================================
- * Custom psutil definitions for modules loaded at runtime.
- * ================================================================
- */
+    PULONG ReturnLength);
 
-_NtQuerySystemInformation \
-    psutil_NtQuerySystemInformation;
-
-_NtQueryInformationProcess \
-    psutil_NtQueryInformationProcess;
-
-_NtSetInformationProcess
-    psutil_NtSetInformationProcess;
-
-_WinStationQueryInformationW \
-    psutil_WinStationQueryInformationW;
-
-_RtlIpv4AddressToStringA \
-    psutil_rtlIpv4AddressToStringA;
-
-_RtlIpv6AddressToStringA \
-    psutil_rtlIpv6AddressToStringA;
-
-_GetExtendedTcpTable \
-    psutil_GetExtendedTcpTable;
-
-_GetExtendedUdpTable \
-    psutil_GetExtendedUdpTable;
-
-_GetActiveProcessorCount \
-    psutil_GetActiveProcessorCount;
-
-_GetTickCount64 \
-    psutil_GetTickCount64;
-
-_NtQueryObject \
-    psutil_NtQueryObject;
-
-_GetLogicalProcessorInformationEx \
-    psutil_GetLogicalProcessorInformationEx;
+typedef NTSTATUS (WINAPI *_RtlGetVersion) (
+    PRTL_OSVERSIONINFOW lpVersionInformation
+);
 
 #endif // __NTEXTAPI_H__
diff --git a/psutil/arch/windows/process_handles.c b/psutil/arch/windows/process_handles.c
index f295f18d..6f133ef9 100644
--- a/psutil/arch/windows/process_handles.c
+++ b/psutil/arch/windows/process_handles.c
@@ -502,16 +502,9 @@ cleanup:
  */
 PyObject *
 psutil_get_open_files(long dwPid, HANDLE hProcess) {
-    OSVERSIONINFO osvi;
-
-    ZeroMemory(&osvi, sizeof(OSVERSIONINFO));
-    osvi.dwOSVersionInfoSize = sizeof(OSVERSIONINFO);
-    GetVersionEx(&osvi);
-
     // Threaded version only works for Vista+
-    if (osvi.dwMajorVersion >= 6)
+    if (PSUTIL_WINVER >= PSUTIL_WINDOWS_VISTA)
         return psutil_get_open_files_ntqueryobject(dwPid, hProcess);
     else
         return psutil_get_open_files_getmappedfilename(dwPid, hProcess);
 }
-
diff --git a/psutil/tests/__init__.py b/psutil/tests/__init__.py
index c6fc8cc2..b1c3c7b3 100644
--- a/psutil/tests/__init__.py
+++ b/psutil/tests/__init__.py
@@ -321,8 +321,11 @@ def get_test_subprocess(cmd=None, **kwds):
     kwds.setdefault("cwd", os.getcwd())
     kwds.setdefault("env", os.environ)
     if WINDOWS:
-        # Prevents the subprocess to open error dialogs.
-        kwds.setdefault("creationflags", 0x8000000)  # CREATE_NO_WINDOW
+        # Prevents the subprocess to open error dialogs. This will also
+        # cause stderr to be suppressed, which is suboptimal in order
+        # to debug broken tests.
+        CREATE_NO_WINDOW = 0x8000000
+        kwds.setdefault("creationflags", CREATE_NO_WINDOW)
     if cmd is None:
         safe_rmpath(_TESTFN)
         pyline = "from time import sleep;" \
diff --git a/psutil/tests/test_connections.py b/psutil/tests/test_connections.py
index 7f59a74c..c97fc409 100755
--- a/psutil/tests/test_connections.py
+++ b/psutil/tests/test_connections.py
@@ -474,7 +474,7 @@ class TestSystemWideConnections(Base, unittest.TestCase):
                 import time, os
                 from psutil.tests import create_sockets
                 with create_sockets():
-                    with open('%s', 'w') as f:
+                    with open(r'%s', 'w') as f:
                         f.write(str(os.getpid()))
                     time.sleep(60)
                 """ % fname)