[psutil] r1260 committed - FIx #52: implement process env vars on Windows.

[psu...@googlecode.com]

Revision: 1260
Author: g.rodola
Date: Wed Feb 1 11:05:15 2012
Log: FIx #52: implement process env vars on Windows.
http://code.google.com/p/psutil/source/detail?r=1260

Modified:
/trunk/psutil/_psmswindows.py
/trunk/psutil/_psutil_mswindows.c
/trunk/psutil/_psutil_mswindows.h
/trunk/test/test_memory_leaks.py

=======================================
--- /trunk/psutil/_psmswindows.py Wed Jan 4 14:17:29 2012
+++ /trunk/psutil/_psmswindows.py Wed Feb 1 11:05:15 2012
@@ -297,6 +297,24 @@
else:
return STATUS_RUNNING

+ @wrap_exceptions
+ def get_process_environ(self):
+ if self.pid in (0, 4) or self.pid == 8 and _WIN2000:
+ raise AccessDenied(self.pid, self._process_name)
+ rawstr = _psutil_mswindows.get_process_environ(self.pid)
+ ret = {}
+ values = rawstr.split('\0\0')[0].split('\0')
+ for line in values:
+ pos = line.find('=')
+ if pos:
+ # os.environ forces upper keys; let's do the same
+ key = line[:pos].upper()
+ if key:
+ value = line[pos+1:]
+ assert value, line
+ ret[key] = value
+ return ret
+
@wrap_exceptions
def get_process_cpu_affinity(self):
from_bitmask = lambda x: [i for i in xrange(64) if (1 << i) & x]
=======================================
--- /trunk/psutil/_psutil_mswindows.c Tue Jan 31 10:37:58 2012
+++ /trunk/psutil/_psutil_mswindows.c Wed Feb 1 11:05:15 2012
@@ -2009,6 +2009,104 @@

return Py_BuildValue("s", type_str);
}
+
+
+static PyObject*
+get_process_environ(PyObject* self, PyObject* args)
+{
+ long pid;
+ HANDLE processHandle;
+ PVOID pebAddress;
+ PVOID rtlUserProcParamsAddress;
+ PVOID currentEnv;
+ MEMORY_BASIC_INFORMATION info;
+ WCHAR buf[4096];
+ DWORD size = sizeof(buf);
+ PyObject *returnPyObj = NULL;
+ PyObject *env_from_wchar = NULL;
+ PyObject *env = NULL;
+
+ if (! PyArg_ParseTuple(args, "l", &pid)) {
+ return NULL;
+ }
+
+ processHandle = handle_from_pid(pid);
+ if (processHandle == NULL) {
+ return NULL;
+ }
+
+ pebAddress = GetPebAddress(processHandle);
+
+ // get the address of ProcessParameters
+ if (!ReadProcessMemory(processHandle, (PCHAR)pebAddress + 0x10,
+ &rtlUserProcParamsAddress, sizeof(PVOID), NULL))
+ {
+ CloseHandle(processHandle);
+
+ if (GetLastError() == ERROR_PARTIAL_COPY) {
+ // this occurs quite often with system processes
+ return AccessDenied();
+ }
+ else {
+ return PyErr_SetFromWindowsErr(0);
+ }
+ }
+
+ // 0x48 refers to "Environ" of RTL_USER_PROCESS_PARAMETERS
+ // Different but related:
+ //
http://wj32.wordpress.com/2009/01/24/howto-get-the-command-line-of-processes/
+ if (!ReadProcessMemory(processHandle, (PCHAR)rtlUserProcParamsAddress
+ 0x48,
+ &currentEnv, sizeof(currentEnv), NULL))
+ {
+ CloseHandle(processHandle);
+ if (GetLastError() == ERROR_PARTIAL_COPY) {
+ // this occurs quite often with system processes
+ return AccessDenied();
+ }
+ else {
+ return PyErr_SetFromWindowsErr(0);
+ }
+ }
+
+ memset(buf, '\0', size);
+ size -= 2;
+
+ if (VirtualQueryEx(processHandle, currentEnv, &info, sizeof(info))) {
+ if (size > info.RegionSize) {
+ size = info.RegionSize;
+ }
+ }
+
+ if (!ReadProcessMemory(processHandle, currentEnv, buf, size, NULL)) {
+ CloseHandle(processHandle);
+ free(buf);
+ if (GetLastError() == ERROR_PARTIAL_COPY) {
+ // this occurs quite often with system processes or when
+ // putenv() was used in case of current process
+ return AccessDenied();
+ }
+ else {
+ return PyErr_SetFromWindowsErr(0);
+ }
+ }
+
+ env_from_wchar = PyUnicode_FromWideChar(buf, size);
+
+ #if PY_MAJOR_VERSION >= 3
+ env = PyUnicode_FromObject(env_from_wchar);
+ #else
+ env = PyUnicode_AsUTF8String(env_from_wchar);
+ #endif
+
+ // decrement the reference count on our temp unicode str to avoid
+ // mem leak
+ Py_XDECREF(env_from_wchar);
+ returnPyObj = Py_BuildValue("N", env);
+
+ CloseHandle(processHandle);
+ free(buf);
+ return returnPyObj;
+}

// ------------------------ Python init ---------------------------
@@ -2063,6 +2161,8 @@
"Get process I/O counters."},
{"is_process_suspended", is_process_suspended, METH_VARARGS,
"Return True if one of the process threads is in a suspended
state"},
+ {"get_process_environ", get_process_environ, METH_VARARGS,
+ "Return process env vars as a string."},

// --- system-related functions

=======================================
--- /trunk/psutil/_psutil_mswindows.h Sun Dec 25 06:55:26 2011
+++ /trunk/psutil/_psutil_mswindows.h Wed Feb 1 11:05:15 2012
@@ -35,6 +35,7 @@
static PyObject* is_process_suspended(PyObject* self, PyObject* args);
static PyObject* get_process_cpu_affinity(PyObject* self, PyObject* args);
static PyObject* set_process_cpu_affinity(PyObject* self, PyObject* args);
+static PyObject* get_process_environ(PyObject* self, PyObject* args);

// --- system-related functions

=======================================
--- /trunk/test/test_memory_leaks.py Wed Jan 4 14:17:29 2012
+++ /trunk/test/test_memory_leaks.py Wed Feb 1 11:05:15 2012
@@ -135,6 +135,9 @@
def test_get_memory_info(self):
self.execute('get_memory_info')

+ def test_get_environ(self):
+ self.execute('get_environ')
+
def test_is_running(self):
self.execute('is_running')