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

10 files changed, 116 insertions, 112 deletions

diff --git a/DEVGUIDE.rst b/DEVGUIDE.rst
index 1d1baf1f..df031bde 100644
--- a/DEVGUIDE.rst
+++ b/DEVGUIDE.rst
@@ -82,7 +82,7 @@ On Windows:
 
 .. code-block:: bat
 
-    set TSCRIPT=foo.py && make test
+    make test foo.py
 
 Adding a new feature
 ====================
diff --git a/HISTORY.rst b/HISTORY.rst
index 6b2a23f7..537b8b29 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -33,9 +33,6 @@ XXXX-XX-XX
   a 64-bit process in 32-bit-WoW mode. Now it raises AccessDenied.
 - 1427_: [OSX] Process cmdline() and environ() may erroneously raise OSError
   on failed malloc().
-- 1431_: [Windows] GetNativeSystemInfo is not used instead of GetSystemInfo in
-  order to support WoW64 processes. Affected APIs are psutil.cpu_count(),
-  and Process memory_maps() and memory_info_exe() ("uss" field).
 - 1432_: [Windows] Process.memory_info_ex()'s USS memory is miscalculated
   because we're not using the actual system PAGESIZE.
 - 1439_: [NetBSD] Process.connections() may return incomplete results if using
diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
index e9db62a5..5c05ac38 100644
--- a/psutil/_psutil_windows.c
+++ b/psutil/_psutil_windows.c
@@ -106,12 +106,12 @@ psutil_get_num_cpus(int fail_on_err) {
     }
     else {
         psutil_debug("GetActiveProcessorCount() not available; "
-                     "using GetNativeSystemInfo()");
+                     "using GetSystemInfo()");
         ncpus = (unsigned int)PSUTIL_SYSTEM_INFO.dwNumberOfProcessors;
-        if ((ncpus == 0) && (fail_on_err == 1)) {
+        if ((ncpus <= 0) && (fail_on_err == 1)) {
             PyErr_SetString(
                 PyExc_RuntimeError,
-                "GetNativeSystemInfo() failed to retrieve CPU count");
+                "GetSystemInfo() failed to retrieve CPU count");
         }
     }
     return ncpus;
@@ -584,13 +584,13 @@ psutil_proc_cmdline(PyObject *self, PyObject *args, PyObject *kwdict) {
     if ((pid == 0) || (pid == 4))
         return Py_BuildValue("[]");
 
-    use_peb = (py_usepeb == Py_True);
     pid_return = psutil_pid_is_running(pid);
     if (pid_return == 0)
         return NoSuchProcess("");
     if (pid_return == -1)
         return NULL;
 
+    use_peb = (py_usepeb == Py_True) ? 1 : 0;
     return psutil_get_cmdline(pid, use_peb);
 }
 
@@ -2186,7 +2186,7 @@ psutil_disk_io_counters(PyObject *self, PyObject *args) {
     char szDeviceDisplay[MAX_PATH];
     int devNum;
     int i;
-    size_t ioctrlSize;
+    DWORD ioctrlSize;
     BOOL ret;
     PyObject *py_retdict = PyDict_New();
     PyObject *py_tuple = NULL;
@@ -2747,11 +2747,7 @@ static char *get_region_protection_string(ULONG protection) {
  */
 static PyObject *
 psutil_proc_memory_maps(PyObject *self, PyObject *args) {
-#ifdef _WIN64
-    MEMORY_BASIC_INFORMATION64 basicInfo;
-#else
     MEMORY_BASIC_INFORMATION basicInfo;
-#endif
     DWORD pid;
     HANDLE hProcess = NULL;
     PVOID baseAddress;
@@ -3297,7 +3293,7 @@ static PyObject *
 psutil_cpu_freq(PyObject *self, PyObject *args) {
     PROCESSOR_POWER_INFORMATION *ppi;
     NTSTATUS ret;
-    size_t size;
+    ULONG size;
     LPBYTE pBuffer = NULL;
     ULONG current;
     ULONG max;
diff --git a/psutil/_pswindows.py b/psutil/_pswindows.py
index 0ab8afe4..3fcee450 100644
--- a/psutil/_pswindows.py
+++ b/psutil/_pswindows.py
@@ -731,13 +731,18 @@ class Process(object):
 
     @wrap_exceptions
     def cmdline(self):
-        try:
+        if cext.WINVER >= cext.WINDOWS_8_1:
+            # PEB method detects cmdline changes but requires more
+            # privileges: https://github.com/giampaolo/psutil/pull/1398
+            try:
+                ret = cext.proc_cmdline(self.pid, use_peb=True)
+            except OSError as err:
+                if err.errno in ACCESS_DENIED_ERRSET:
+                    ret = cext.proc_cmdline(self.pid, use_peb=False)
+                else:
+                    raise
+        else:
             ret = cext.proc_cmdline(self.pid, use_peb=True)
-        except OSError as err:
-            if err.errno in ACCESS_DENIED_ERRSET:
-                ret = cext.proc_cmdline(self.pid, use_peb=False)
-            else:
-                raise
         if PY3:
             return ret
         else:
diff --git a/psutil/arch/windows/global.c b/psutil/arch/windows/global.c
index 6134687d..a622b635 100644
--- a/psutil/arch/windows/global.c
+++ b/psutil/arch/windows/global.c
@@ -180,7 +180,7 @@ psutil_set_winver() {
 
 static int
 psutil_load_sysinfo() {
-    GetNativeSystemInfo(&PSUTIL_SYSTEM_INFO);
+    GetSystemInfo(&PSUTIL_SYSTEM_INFO);
     return 0;
 }
 
diff --git a/psutil/arch/windows/ntextapi.h b/psutil/arch/windows/ntextapi.h
index 8007ec01..3927edf3 100644
--- a/psutil/arch/windows/ntextapi.h
+++ b/psutil/arch/windows/ntextapi.h
@@ -185,44 +185,46 @@ typedef enum _KTHREAD_STATE {
 } KTHREAD_STATE, *PKTHREAD_STATE;
 
 typedef enum _KWAIT_REASON {
-    Executive = 0,
-    FreePage = 1,
-    PageIn = 2,
-    PoolAllocation = 3,
-    DelayExecution = 4,
-    Suspended = 5,
-    UserRequest = 6,
-    WrExecutive = 7,
-    WrFreePage = 8,
-    WrPageIn = 9,
-    WrPoolAllocation = 10,
-    WrDelayExecution = 11,
-    WrSuspended = 12,
-    WrUserRequest = 13,
-    WrEventPair = 14,
-    WrQueue = 15,
-    WrLpcReceive = 16,
-    WrLpcReply = 17,
-    WrVirtualMemory = 18,
-    WrPageOut = 19,
-    WrRendezvous = 20,
-    Spare2 = 21,
-    Spare3 = 22,
-    Spare4 = 23,
-    Spare5 = 24,
-    WrCalloutStack = 25,
-    WrKernel = 26,
-    WrResource = 27,
-    WrPushLock = 28,
-    WrMutex = 29,
-    WrQuantumEnd = 30,
-    WrDispatchInt = 31,
-    WrPreempted = 32,
-    WrYieldExecution = 33,
-    WrFastMutex = 34,
-    WrGuardedMutex = 35,
-    WrRundown = 36,
-    MaximumWaitReason = 37
+    Executive,
+    FreePage,
+    PageIn,
+    PoolAllocation,
+    DelayExecution,
+    Suspended,
+    UserRequest,
+    WrExecutive,
+    WrFreePage,
+    WrPageIn,
+    WrPoolAllocation,
+    WrDelayExecution,
+    WrSuspended,
+    WrUserRequest,
+    WrEventPair,
+    WrQueue,
+    WrLpcReceive,
+    WrLpcReply,
+    WrVirtualMemory,
+    WrPageOut,
+    WrRendezvous,
+    WrKeyedEvent,
+    WrTerminated,
+    WrProcessInSwap,
+    WrCpuRateControl,
+    WrCalloutStack,
+    WrKernel,
+    WrResource,
+    WrPushLock,
+    WrMutex,
+    WrQuantumEnd,
+    WrDispatchInt,
+    WrPreempted,
+    WrYieldExecution,
+    WrFastMutex,
+    WrGuardedMutex,
+    WrRundown,
+    WrAlertByThreadId,
+    WrDeferredPreempt,
+    MaximumWaitReason
 } KWAIT_REASON, *PKWAIT_REASON;
 
 typedef struct _SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX {
diff --git a/psutil/arch/windows/process_handles.c b/psutil/arch/windows/process_handles.c
index 6f133ef9..8b899972 100644
--- a/psutil/arch/windows/process_handles.c
+++ b/psutil/arch/windows/process_handles.c
@@ -169,11 +169,11 @@ psutil_get_open_files_ntqueryobject(long dwPid, HANDLE hProcess) {
         hHandle = &pHandleInfo->Handles[i];
 
         // Check if this hHandle belongs to the PID the user specified.
-        if (hHandle->UniqueProcessId != (HANDLE)dwPid)
+        if (hHandle->UniqueProcessId != (ULONG_PTR)dwPid)
             goto loop_cleanup;
 
         if (!DuplicateHandle(hProcess,
-                             hHandle->HandleValue,
+                             (HANDLE)hHandle->HandleValue,
                              GetCurrentProcess(),
                              &g_hFile,
                              0,
@@ -318,7 +318,7 @@ psutil_get_open_files_getmappedfilename(long dwPid, HANDLE hProcess) {
     PyObject*                           py_path = NULL;
     ULONG                               dwSize = 0;
     LPVOID                              pMem = NULL;
-    TCHAR                               pszFilename[MAX_PATH+1];
+    wchar_t                             pszFilename[MAX_PATH+1];
 
     if (g_initialized == FALSE)
         psutil_get_open_files_init(FALSE);
@@ -365,11 +365,11 @@ psutil_get_open_files_getmappedfilename(long dwPid, HANDLE hProcess) {
         hHandle = &pHandleInfo->Handles[i];
 
         // Check if this hHandle belongs to the PID the user specified.
-        if (hHandle->UniqueProcessId != (HANDLE)dwPid)
+        if (hHandle->UniqueProcessId != (ULONG_PTR)dwPid)
             goto loop_cleanup;
 
         if (!DuplicateHandle(hProcess,
-                             hHandle->HandleValue,
+                             (HANDLE)hHandle->HandleValue,
                              GetCurrentProcess(),
                              &hFile,
                              0,
@@ -409,7 +409,7 @@ psutil_get_open_files_getmappedfilename(long dwPid, HANDLE hProcess) {
         }
 
         dwSize = GetMappedFileName(
-            GetCurrentProcess(), pMem, pszFilename, MAX_PATH);
+            GetCurrentProcess(), pMem, (LPSTR)pszFilename, MAX_PATH);
         if (dwSize == 0) {
             /*
             printf("[%d] GetMappedFileName (%#x): %#x \n",
diff --git a/psutil/arch/windows/process_info.c b/psutil/arch/windows/process_info.c
index 3966bc71..946a01cb 100644
--- a/psutil/arch/windows/process_info.c
+++ b/psutil/arch/windows/process_info.c
@@ -735,24 +735,24 @@ error:
 
 
 /*
- * Get process cmdline() by using NtQueryInformationProcess. This is
- * useful on Windows 8.1+ in order to get less ERROR_ACCESS_DENIED
- * errors when querying privileged PIDs.
+ * Get process cmdline by using NtQueryInformationProcess. This is a
+ * method alternative to PEB which is less likely to result in
+ * AccessDenied. Requires Windows 8.1+.
  */
 static int
-psutil_get_cmdline_data(long pid, WCHAR **pdata, SIZE_T *psize) {
+psutil_cmdline_query_proc(long pid, WCHAR **pdata, SIZE_T *psize) {
     HANDLE hProcess;
-    ULONG ret_length = 4096;
+    ULONG bufLen = 0;
     NTSTATUS status;
-    char * cmdline_buffer = NULL;
-    WCHAR * cmdline_buffer_wchar = NULL;
+    char * buffer = NULL;
+    WCHAR * bufWchar = NULL;
     PUNICODE_STRING tmp = NULL;
-    DWORD string_size;
+    size_t size;
     int ProcessCommandLineInformation = 60;
 
-    cmdline_buffer = calloc(ret_length, 1);
-    if (cmdline_buffer == NULL) {
-        PyErr_NoMemory();
+    if (PSUTIL_WINVER < PSUTIL_WINDOWS_8_1) {
+        PyErr_SetString(
+            PyExc_RuntimeError, "requires Windows 8.1+");
         goto error;
     }
 
@@ -766,7 +766,7 @@ psutil_get_cmdline_data(long pid, WCHAR **pdata, SIZE_T *psize) {
         ProcessCommandLineInformation,
         NULL,
         0,
-        &ret_length);
+        &bufLen);
     if (status != STATUS_BUFFER_OVERFLOW && \
             status != STATUS_BUFFER_TOO_SMALL && \
             status != STATUS_INFO_LENGTH_MISMATCH) {
@@ -774,13 +774,20 @@ psutil_get_cmdline_data(long pid, WCHAR **pdata, SIZE_T *psize) {
         goto error;
     }
 
+    // allocate memory
+    buffer = calloc(bufLen, 1);
+    if (buffer == NULL) {
+        PyErr_NoMemory();
+        goto error;
+    }
+
     // get the cmdline
     status = psutil_NtQueryInformationProcess(
         hProcess,
         ProcessCommandLineInformation,
-        cmdline_buffer,
-        ret_length,
-        &ret_length
+        buffer,
+        bufLen,
+        &bufLen
     );
     if (! NT_SUCCESS(status)) {
         PyErr_SetFromOSErrnoWithSyscall("NtQueryInformationProcess(withlen)");
@@ -788,23 +795,23 @@ psutil_get_cmdline_data(long pid, WCHAR **pdata, SIZE_T *psize) {
     }
 
     // build the string
-    tmp = (PUNICODE_STRING)cmdline_buffer;
-    string_size = wcslen(tmp->Buffer) + 1;
-    cmdline_buffer_wchar = (WCHAR *)calloc(string_size, sizeof(WCHAR));
-    if (cmdline_buffer_wchar == NULL) {
+    tmp = (PUNICODE_STRING)buffer;
+    size = wcslen(tmp->Buffer) + 1;
+    bufWchar = (WCHAR *)calloc(size, sizeof(WCHAR));
+    if (bufWchar == NULL) {
         PyErr_NoMemory();
         goto error;
     }
-    wcscpy_s(cmdline_buffer_wchar, string_size, tmp->Buffer);
-    *pdata = cmdline_buffer_wchar;
-    *psize = string_size * sizeof(WCHAR);
-    free(cmdline_buffer);
+    wcscpy_s(bufWchar, size, tmp->Buffer);
+    *pdata = bufWchar;
+    *psize = size * sizeof(WCHAR);
+    free(buffer);
     CloseHandle(hProcess);
     return 0;
 
 error:
-    if (cmdline_buffer != NULL)
-        free(cmdline_buffer);
+    if (buffer != NULL)
+        free(buffer);
     if (hProcess != NULL)
         CloseHandle(hProcess);
     return -1;
@@ -827,24 +834,18 @@ psutil_get_cmdline(long pid, int use_peb) {
     int func_ret;
 
     /*
-    By defaut, still use PEB (if command line params have been patched in
-    the PEB, we will get the actual ones). Reading the PEB to get the
-    command line parameters still seem to be the best method if somebody
-    has tampered with the parameters after creating the process.
+    Reading the PEB to get the cmdline seem to be the best method if
+    somebody has tampered with the parameters after creating the process.
     For instance, create a process as suspended, patch the command line
-    in its PEB and unfreeze it.
-    The process will use the "new" parameters whereas the system
-    (with NtQueryInformationProcess) will give you the "old" ones
-    See:
+    in its PEB and unfreeze it. It requires more privileges than
+    NtQueryInformationProcess though (the fallback):
     - https://github.com/giampaolo/psutil/pull/1398
     - https://blog.xpnsec.com/how-to-argue-like-cobalt-strike/
     */
-    if (use_peb == 1) {
+    if (use_peb == 1)
         func_ret = psutil_get_process_data(pid, KIND_CMDLINE, &data, &size);
-    }
-    else {
-        func_ret = psutil_get_cmdline_data(pid, &data, &size);
-    }
+    else
+        func_ret = psutil_cmdline_query_proc(pid, &data, &size);
     if (func_ret != 0)
         goto out;
 
diff --git a/psutil/tests/test_windows.py b/psutil/tests/test_windows.py
index ff63cd79..c98d892c 100755
--- a/psutil/tests/test_windows.py
+++ b/psutil/tests/test_windows.py
@@ -82,7 +82,7 @@ class TestCpuAPIs(unittest.TestCase):
     def test_cpu_count_vs_GetSystemInfo(self):
         # Will likely fail on many-cores systems:
         # https://stackoverflow.com/questions/31209256
-        sys_value = win32api.GetNativeSystemInfo()[5]
+        sys_value = win32api.GetSystemInfo()[5]
         psutil_value = psutil.cpu_count()
         self.assertEqual(sys_value, psutil_value)
 
diff --git a/scripts/internal/winmake.py b/scripts/internal/winmake.py
index cd26c67e..c35853c5 100755
--- a/scripts/internal/winmake.py
+++ b/scripts/internal/winmake.py
@@ -29,7 +29,7 @@ if APPVEYOR:
     PYTHON = sys.executable
 else:
     PYTHON = os.getenv('PYTHON', sys.executable)
-TSCRIPT = os.getenv('TSCRIPT', 'psutil\\tests\\__main__.py')
+TEST_SCRIPT = 'psutil\\tests\\__main__.py'
 GET_PIP_URL = "https://bootstrap.pypa.io/get-pip.py"
 PY3 = sys.version_info[0] == 3
 HERE = os.path.abspath(os.path.dirname(__file__))
@@ -350,9 +350,16 @@ def flake8():
 @cmd
 def test():
     """Run tests"""
+    try:
+        arg = sys.argv[2]
+    except IndexError:
+        arg = TEST_SCRIPT
+
     install()
     test_setup()
-    sh("%s %s" % (PYTHON, TSCRIPT))
+    cmdline = "%s %s" % (PYTHON, arg)
+    safe_print(cmdline)
+    sh(cmdline)
 
 
 @cmd
@@ -361,7 +368,7 @@ def coverage():
     # Note: coverage options are controlled by .coveragerc file
     install()
     test_setup()
-    sh("%s -m coverage run %s" % (PYTHON, TSCRIPT))
+    sh("%s -m coverage run %s" % (PYTHON, TEST_SCRIPT))
     sh("%s -m coverage report" % PYTHON)
     sh("%s -m coverage html" % PYTHON)
     sh("%s -m webbrowser -t htmlcov/index.html" % PYTHON)
@@ -426,11 +433,7 @@ def test_contracts():
 @cmd
 def test_by_name():
     """Run test by name"""
-    try:
-        safe_print(sys.argv)
-        name = sys.argv[2]
-    except IndexError:
-        sys.exit('second arg missing')
+    name = sys.argv[2]
     install()
     test_setup()
     sh("%s -m unittest -v %s" % (PYTHON, name))