1 files changed, 189 insertions, 0 deletions
diff --git a/psutil/arch/windows/proc_utils.c b/psutil/arch/windows/proc_utils.c
new file mode 100644
index 00000000..dac1129c
--- /dev/null
+++ b/psutil/arch/windows/proc_utils.c
@@ -0,0 +1,189 @@
+/*
+ * Copyright (c) 2009, Jay Loden, 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.
+ *
+ * Helper process functions.
+ */
+
+#include <Python.h>
+#include <windows.h>
+#include <Psapi.h>  // EnumProcesses
+
+#include "../../_psutil_common.h"
+#include "proc_utils.h"
+
+
+DWORD *
+psutil_get_pids(DWORD *numberOfReturnedPIDs) {
+    // Win32 SDK says the only way to know if our process array
+    // wasn't large enough is to check the returned size and make
+    // sure that it doesn't match the size of the array.
+    // If it does we allocate a larger array and try again
+
+    // Stores the actual array
+    DWORD *procArray = NULL;
+    DWORD procArrayByteSz;
+    int procArraySz = 0;
+
+    // Stores the byte size of the returned array from enumprocesses
+    DWORD enumReturnSz = 0;
+
+    do {
+        procArraySz += 1024;
+        if (procArray != NULL)
+            free(procArray);
+        procArrayByteSz = procArraySz * sizeof(DWORD);
+        procArray = malloc(procArrayByteSz);
+        if (procArray == NULL) {
+            PyErr_NoMemory();
+            return NULL;
+        }
+        if (! EnumProcesses(procArray, procArrayByteSz, &enumReturnSz)) {
+            free(procArray);
+            PyErr_SetFromWindowsErr(0);
+            return NULL;
+        }
+    } while (enumReturnSz == procArraySz * sizeof(DWORD));
+
+    // The number of elements is the returned size / size of each element
+    *numberOfReturnedPIDs = enumReturnSz / sizeof(DWORD);
+
+    return procArray;
+}
+
+
+// Return 1 if PID exists, 0 if not, -1 on error.
+int
+psutil_pid_in_pids(DWORD pid) {
+    DWORD *proclist = NULL;
+    DWORD numberOfReturnedPIDs;
+    DWORD i;
+
+    proclist = psutil_get_pids(&numberOfReturnedPIDs);
+    if (proclist == NULL) {
+        psutil_debug("psutil_get_pids() failed");
+        return -1;
+    }
+    for (i = 0; i < numberOfReturnedPIDs; i++) {
+        if (proclist[i] == pid) {
+            free(proclist);
+            return 1;
+        }
+    }
+    free(proclist);
+    return 0;
+}
+
+
+// Given a process handle checks whether it's actually running. If it
+// does return the handle, else return NULL with Python exception set.
+// This is needed because OpenProcess API sucks.
+HANDLE
+psutil_check_phandle(HANDLE hProcess, DWORD pid, int check_exit_code) {
+    DWORD exitCode;
+
+    if (hProcess == NULL) {
+        if (GetLastError() == ERROR_INVALID_PARAMETER) {
+            // Yeah, this is the actual error code in case of
+            // "no such process".
+            NoSuchProcess("OpenProcess -> ERROR_INVALID_PARAMETER");
+            return NULL;
+        }
+        if (GetLastError() == ERROR_SUCCESS) {
+            // Yeah, it's this bad.
+            // https://github.com/giampaolo/psutil/issues/1877
+            if (psutil_pid_in_pids(pid) == 1) {
+                psutil_debug("OpenProcess -> ERROR_SUCCESS turned into AD");
+                AccessDenied("OpenProcess -> ERROR_SUCCESS");
+            }
+            else {
+                psutil_debug("OpenProcess -> ERROR_SUCCESS turned into NSP");
+                NoSuchProcess("OpenProcess -> ERROR_SUCCESS");
+            }
+            return NULL;
+        }
+        PyErr_SetFromOSErrnoWithSyscall("OpenProcess");
+        return NULL;
+    }
+
+    if (check_exit_code == 0)
+        return hProcess;
+
+    if (GetExitCodeProcess(hProcess, &exitCode)) {
+        // XXX - maybe STILL_ACTIVE is not fully reliable as per:
+        // http://stackoverflow.com/questions/1591342/#comment47830782_1591379
+        if (exitCode == STILL_ACTIVE) {
+            return hProcess;
+        }
+        if (psutil_pid_in_pids(pid) == 1) {
+            return hProcess;
+        }
+        CloseHandle(hProcess);
+        NoSuchProcess("GetExitCodeProcess != STILL_ACTIVE");
+        return NULL;
+    }
+
+    if (GetLastError() == ERROR_ACCESS_DENIED) {
+        psutil_debug("GetExitCodeProcess -> ERROR_ACCESS_DENIED (ignored)");
+        SetLastError(0);
+        return hProcess;
+    }
+    PyErr_SetFromOSErrnoWithSyscall("GetExitCodeProcess");
+    CloseHandle(hProcess);
+    return NULL;
+}
+
+
+// A wrapper around OpenProcess setting NSP exception if process no
+// longer exists. *pid* is the process PID, *access* is the first
+// argument to OpenProcess.
+// Return a process handle or NULL with exception set.
+HANDLE
+psutil_handle_from_pid(DWORD pid, DWORD access) {
+    HANDLE hProcess;
+
+    if (pid == 0) {
+        // otherwise we'd get NoSuchProcess
+        return AccessDenied("automatically set for PID 0");
+    }
+
+    hProcess = OpenProcess(access, FALSE, pid);
+
+    if ((hProcess == NULL) && (GetLastError() == ERROR_ACCESS_DENIED)) {
+        PyErr_SetFromOSErrnoWithSyscall("OpenProcess");
+        return NULL;
+    }
+
+    hProcess = psutil_check_phandle(hProcess, pid, 1);
+    return hProcess;
+}
+
+
+// Check for PID existence. Return 1 if pid exists, 0 if not, -1 on error.
+int
+psutil_pid_is_running(DWORD pid) {
+    HANDLE hProcess;
+
+    // Special case for PID 0 System Idle Process
+    if (pid == 0)
+        return 1;
+    if (pid < 0)
+        return 0;
+
+    hProcess = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, pid);
+
+    // Access denied means there's a process to deny access to.
+    if ((hProcess == NULL) && (GetLastError() == ERROR_ACCESS_DENIED))
+        return 1;
+
+    hProcess = psutil_check_phandle(hProcess, pid, 1);
+    if (hProcess != NULL) {
+        CloseHandle(hProcess);
+        return 1;
+    }
+
+    CloseHandle(hProcess);
+    PyErr_Clear();
+    return psutil_pid_in_pids(pid);
+}