#799, oneshot(), windows: expose C functions to OpenProcess and Close…

…Handle in order to keep the handle reference at Python level and allow caching

psutil/_psutil_windows.c

@@ -2661,19 +2661,14 @@ psutil_users(PyObject *self, PyObject *args) {
 static PyObject *
 psutil_proc_num_handles(PyObject *self, PyObject *args) {
     DWORD pid;
-    HANDLE hProcess;
+    unsigned long handle;
     DWORD handleCount;
 
-    if (! PyArg_ParseTuple(args, "l", &pid))
+    if (! PyArg_ParseTuple(args, "lk", &pid, &handle))
         return NULL;
-    hProcess = psutil_handle_from_pid(pid);
-    if (NULL == hProcess)
-        return NULL;
-    if (! GetProcessHandleCount(hProcess, &handleCount)) {
-        CloseHandle(hProcess);
+    if (! GetProcessHandleCount((HANDLE)handle, &handleCount)) {
         return PyErr_SetFromWindowsErr(0);
     }
-    CloseHandle(hProcess);
     return Py_BuildValue("k", handleCount);
 }
 
@@ -3430,6 +3425,10 @@ PsutilMethods[] = {
     // --- windows API bindings
     {"win32_QueryDosDevice", psutil_win32_QueryDosDevice, METH_VARARGS,
      "QueryDosDevice binding"},
+    {"win32_OpenProcess", psutil_win32_OpenProcess, METH_VARARGS,
+     "Given a PID return a Python int which points to a process handle."},
+    {"win32_CloseHandle", psutil_win32_CloseHandle, METH_VARARGS,
+     "Given a Python int referencing a process handle it close the handle."},
 
     {NULL, NULL, 0, NULL}
 };

psutil/_pswindows.py

@@ -546,18 +546,48 @@ def wrapper(self, *args, **kwargs):
 class Process(object):
     """Wrapper class around underlying C implementation."""
 
-    __slots__ = ["pid", "_name", "_ppid"]
+    __slots__ = ["pid", "_name", "_ppid", "_inctx", "_handle"]
 
     def __init__(self, pid):
         self.pid = pid
         self._name = None
         self._ppid = None
+        self._inctx = False
+        self._handle = None
 
     def oneshot_enter(self):
-        pass
+        self._inctx = True
 
     def oneshot_exit(self):
-        pass
+        self._inctx = False
+        if self._handle:
+            cext.win32_CloseHandle(self._handle)
+            self._handle = None
+
+    def get_handle(self):
+        """Get a handle to this process.
+        If we're in oneshot() ctx manager tries to return the
+        cached handle.
+        """
+        if self._inctx:
+            handle = self._handle or cext.win32_OpenProcess(self.pid)
+            return handle
+        else:
+            return cext.win32_OpenProcess(self.pid)
+
+    @contextlib.contextmanager
+    def handle_ctx(self):
+        """Get a handle to this process.
+        If we're not in oneshot() ctx close the handle on exit
+        else tries to return the cached handle and avoid to close
+        the handle (will be close on oneshot() exit).
+        """
+        handle = self.get_handle()
+        try:
+            yield handle
+        finally:
+            if not self._inctx:
+                cext.win32_CloseHandle(handle)
 
     @wrap_exceptions
     def name(self):
@@ -837,7 +867,8 @@ def to_bitmask(l):
     @wrap_exceptions
     def num_handles(self):
         try:
-            return cext.proc_num_handles(self.pid)
+            with self.handle_ctx() as handle:
+                return cext.proc_num_handles(self.pid, handle)
         except OSError as err:
             if err.errno in ACCESS_DENIED_SET:
                 return ntpinfo(*cext.proc_info(self.pid)).num_handles

psutil/arch/windows/process_info.c

@@ -67,6 +67,38 @@ psutil_handle_from_pid(DWORD pid) {
 }
 
 
+/*
+ * Given a PID return a Python int which points to its process handle.
+ */
+PyObject *
+psutil_win32_OpenProcess(PyObject *self, PyObject *args) {
+    HANDLE handle;
+    long pid;
+
+    if (! PyArg_ParseTuple(args, "l", &pid))
+        return NULL;
+    handle = psutil_handle_from_pid(pid);
+    if (handle == NULL)
+        return NULL;
+    return HANDLE_TO_PYNUM(handle);
+}
+
+
+/*
+ * Given a Python int referencing a process handle close the process handle.
+ */
+PyObject *
+psutil_win32_CloseHandle(PyObject *self, PyObject *args) {
+    unsigned long handle;
+
+    if (! PyArg_ParseTuple(args, "k", &handle))
+        return NULL;
+    // TODO: may want to check return value;
+    CloseHandle((HANDLE)handle);
+    Py_RETURN_NONE;
+}
+
+
 DWORD *
 psutil_get_pids(DWORD *numberOfReturnedPIDs) {
     // Win32 SDK says the only way to know if our process array
@@ -615,7 +647,7 @@ static int psutil_get_process_data(long pid,
                 src = procParameters.CommandLine.Buffer;
                 size = procParameters.CommandLine.Length;
                 break;
-            case KIND_CWD: 
+            case KIND_CWD:
                 src = procParameters.CurrentDirectoryPath.Buffer;
                 size = procParameters.CurrentDirectoryPath.Length;
                 break;

psutil/arch/windows/process_info.h

@@ -12,9 +12,16 @@
 #include "security.h"
 #include "ntextapi.h"
 
+#define HANDLE_TO_PYNUM(handle) PyLong_FromUnsignedLong((unsigned long) handle)
+#define PYNUM_TO_HANDLE(obj) ((HANDLE)PyLong_AsUnsignedLong(obj))
+
+
 DWORD* psutil_get_pids(DWORD *numberOfReturnedPIDs);
 HANDLE psutil_handle_from_pid(DWORD pid);
 HANDLE psutil_handle_from_pid_waccess(DWORD pid, DWORD dwDesiredAccess);
+PyObject* psutil_win32_OpenProcess(PyObject *self, PyObject *args);
+PyObject* psutil_win32_CloseHandle(PyObject *self, PyObject *args);
+
 int psutil_handlep_is_running(HANDLE hProcess);
 int psutil_pid_in_proclist(DWORD pid);
 int psutil_pid_is_running(DWORD pid);