implement Process.environ() on Windowsrevive-environ-windows

8 files changed, 183 insertions, 7 deletions

diff --git a/psutil/__init__.py b/psutil/__init__.py
index c034ad6f..8a896632 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -744,7 +744,7 @@ class Process(object):
             else:
                 self._proc.cpu_affinity_set(list(set(cpus)))
 
-    # Linux and OSX only
+    # Linux, OSX and Windows only
     if hasattr(_psplatform.Process, "environ"):
         def environ(self):
             """The environment variables of the process as a dict.  Note: this
diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
index aa80519e..10ff5dc0 100644
--- a/psutil/_psutil_windows.c
+++ b/psutil/_psutil_windows.c
@@ -592,6 +592,29 @@ psutil_proc_cmdline(PyObject *self, PyObject *args) {
 
 
 /*
+ * Return process cmdline as a Python list of cmdline arguments.
+ */
+static PyObject *
+psutil_proc_environ(PyObject *self, PyObject *args) {
+    long pid;
+    int pid_return;
+
+    if (! PyArg_ParseTuple(args, "l", &pid))
+        return NULL;
+    if ((pid == 0) || (pid == 4))
+        return Py_BuildValue("s", "");
+
+    pid_return = psutil_pid_is_running(pid);
+    if (pid_return == 0)
+        return NoSuchProcess();
+    if (pid_return == -1)
+        return NULL;
+
+    return psutil_get_environ(pid);
+}
+
+
+/*
  * Return process executable path.
  */
 static PyObject *
@@ -2991,6 +3014,8 @@ PsutilMethods[] = {
 
     {"proc_cmdline", psutil_proc_cmdline, METH_VARARGS,
      "Return process cmdline as a list of cmdline arguments"},
+    {"proc_environ", psutil_proc_environ, METH_VARARGS,
+     "Return process environment data"},
     {"proc_exe", psutil_proc_exe, METH_VARARGS,
      "Return path of the process executable"},
     {"proc_name", psutil_proc_name, METH_VARARGS,
diff --git a/psutil/_psutil_windows.h b/psutil/_psutil_windows.h
index c77f64e9..0c22bf7e 100644
--- a/psutil/_psutil_windows.h
+++ b/psutil/_psutil_windows.h
@@ -16,6 +16,7 @@ static PyObject* psutil_proc_cpu_times(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_create_time(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_cwd(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_exe(PyObject* self, PyObject* args);
+static PyObject* psutil_proc_environ(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_info(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_io_counters(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_is_suspended(PyObject* self, PyObject* args);
diff --git a/psutil/_pswindows.py b/psutil/_pswindows.py
index 4d39d84f..6392c510 100644
--- a/psutil/_pswindows.py
+++ b/psutil/_pswindows.py
@@ -14,6 +14,7 @@ from . import _common
 from . import _psutil_windows as cext
 from ._common import conn_tmap
 from ._common import isfile_strict
+from ._common import parse_environ_block
 from ._common import sockfam_to_enum
 from ._common import socktype_to_enum
 from ._common import usage_percent
@@ -346,6 +347,10 @@ class Process(object):
         else:
             return [py2_strencode(s) for s in ret]
 
+    @wrap_exceptions
+    def environ(self):
+        return parse_environ_block(cext.proc_environ(self.pid))
+
     def ppid(self):
         try:
             return ppid_map()[self.pid]
diff --git a/psutil/arch/windows/process_info.c b/psutil/arch/windows/process_info.c
index df7ae9fb..ab1f844c 100644
--- a/psutil/arch/windows/process_info.c
+++ b/psutil/arch/windows/process_info.c
@@ -273,6 +273,18 @@ typedef NTSTATUS (NTAPI *_NtWow64ReadVirtualMemory64)(
         IN ULONG64 Size,
         OUT PULONG64 NumberOfBytesRead);
 
+typedef enum {
+    MemoryInformationBasic
+} MEMORY_INFORMATION_CLASS;
+
+typedef NTSTATUS (NTAPI *_NtWow64QueryVirtualMemory64)(
+    IN HANDLE ProcessHandle,
+    IN PVOID64 BaseAddress,
+    IN MEMORY_INFORMATION_CLASS MemoryInformationClass,
+    OUT PMEMORY_BASIC_INFORMATION64 MemoryInformation,
+    IN ULONG64 Size,
+    OUT PULONG64 ReturnLength OPTIONAL);
+
 typedef struct {
     PVOID Reserved1[2];
     PVOID64 PebBaseAddress;
@@ -308,9 +320,66 @@ typedef struct {
 } PEB64;
 #endif
 
+/* Given a pointer into a process's memory, figure out how much data can be
+ * read from it. */
+static int psutil_get_process_region_size(HANDLE hProcess,
+                                          LPCVOID src,
+                                          SIZE_T *psize) {
+    MEMORY_BASIC_INFORMATION info;
+
+    if (!VirtualQueryEx(hProcess, src, &info, sizeof(info))) {
+        PyErr_SetFromWindowsErr(0);
+        return -1;
+    }
+
+    *psize = info.RegionSize - ((char*)src - (char*)info.BaseAddress);
+    return 0;
+}
+
+
+#ifndef _WIN64
+/* Given a pointer into a process's memory, figure out how much data can be
+ * read from it. */
+static int psutil_get_process_region_size64(HANDLE hProcess,
+                                           const PVOID64 src64,
+                                           PULONG64 psize) {
+    static _NtWow64QueryVirtualMemory64 NtWow64QueryVirtualMemory64 = NULL;
+    MEMORY_BASIC_INFORMATION64 info64;
+
+    if (NtWow64QueryVirtualMemory64 == NULL) {
+        NtWow64QueryVirtualMemory64 =
+            (_NtWow64QueryVirtualMemory64)GetProcAddress(
+                    GetModuleHandleA("ntdll.dll"),
+                    "NtWow64QueryVirtualMemory64");
+
+        if (NtWow64QueryVirtualMemory64 == NULL) {
+            PyErr_SetString(PyExc_NotImplementedError,
+                    "NtWow64QueryVirtualMemory64 missing");
+            return -1;
+        }
+    }
+
+    if (!NT_SUCCESS(NtWow64QueryVirtualMemory64(
+                    hProcess,
+                    src64,
+                    0,
+                    &info64,
+                    sizeof(info64),
+                    NULL))) {
+        PyErr_SetFromWindowsErr(0);
+        return -1;
+    }
+
+    *psize = info64.RegionSize - ((char*)src64 - (char*)info64.BaseAddress);
+    return 0;
+}
+#endif
+
+
 enum psutil_process_data_kind {
     KIND_CMDLINE,
     KIND_CWD,
+    KIND_ENVIRON,
 };
 
 /* Get data from the process with the given pid.  The data is returned in the
@@ -321,7 +390,8 @@ enum psutil_process_data_kind {
    is returned, and an appropriate Python exception is set. */
 static int psutil_get_process_data(long pid,
                                    enum psutil_process_data_kind kind,
-                                   WCHAR **pdata) {
+                                   WCHAR **pdata,
+                                   SIZE_T *psize) {
     /* This function is quite complex because there are several cases to be
        considered:
 
@@ -414,6 +484,9 @@ static int psutil_get_process_data(long pid,
                 src = UlongToPtr(procParameters32.CurrentDirectoryPath.Buffer);
                 size = procParameters32.CurrentDirectoryPath.Length;
                 break;
+            case KIND_ENVIRON:
+                src = UlongToPtr(procParameters32.env);
+                break;
         }
     } else
 #else
@@ -433,8 +506,14 @@ static int psutil_get_process_data(long pid,
         if (NtWow64QueryInformationProcess64 == NULL) {
             NtWow64QueryInformationProcess64 =
                 (_NtQueryInformationProcess)GetProcAddress(
-                    GetModuleHandleA("ntdll.dll"),
-                    "NtWow64QueryInformationProcess64");
+                        GetModuleHandleA("ntdll.dll"),
+                        "NtWow64QueryInformationProcess64");
+
+            if (NtWow64QueryInformationProcess64 == NULL) {
+                PyErr_SetString(PyExc_NotImplementedError,
+                                "NtWow64QueryInformationProcess64 missing");
+                goto error;
+            }
         }
 
         if (!NT_SUCCESS(NtWow64QueryInformationProcess64(
@@ -453,6 +532,12 @@ static int psutil_get_process_data(long pid,
                 (_NtWow64ReadVirtualMemory64)GetProcAddress(
                         GetModuleHandleA("ntdll.dll"),
                         "NtWow64ReadVirtualMemory64");
+
+            if (NtWow64ReadVirtualMemory64 == NULL) {
+                PyErr_SetString(PyExc_NotImplementedError,
+                                "NtWow64ReadVirtualMemory64 missing");
+                goto error;
+            }
         }
 
         if (!NT_SUCCESS(NtWow64ReadVirtualMemory64(hProcess,
@@ -483,6 +568,9 @@ static int psutil_get_process_data(long pid,
                 src64 = procParameters64.CurrentDirectoryPath.Buffer,
                 size = procParameters64.CurrentDirectoryPath.Length;
                 break;
+            case KIND_ENVIRON:
+                src64 = procParameters64.env;
+                break;
         }
     } else
 #endif
@@ -531,7 +619,26 @@ static int psutil_get_process_data(long pid,
                 src = procParameters.CurrentDirectoryPath.Buffer;
                 size = procParameters.CurrentDirectoryPath.Length;
                 break;
+            case KIND_ENVIRON:
+                src = procParameters.env;
+                break;
+        }
+    }
+
+    if (kind == KIND_ENVIRON) {
+#ifndef _WIN64
+        if (weAreWow64 && !theyAreWow64) {
+            ULONG64 size64;
+
+            if (psutil_get_process_region_size64(hProcess, src64, &size64) != 0)
+                goto error;
+
+            size = (SIZE_T)size64;
         }
+        else
+#endif
+        if (psutil_get_process_region_size(hProcess, src, &size) != 0)
+            goto error;
     }
 
     buffer = calloc(size + 2, 1);
@@ -559,6 +666,7 @@ static int psutil_get_process_data(long pid,
     }
 
     *pdata = buffer;
+    *psize = size;
 
     return 0;
 
@@ -578,12 +686,13 @@ PyObject *
 psutil_get_cmdline(long pid) {
     PyObject *ret = NULL;
     WCHAR *data = NULL;
+    SIZE_T size;
     PyObject *py_retlist = NULL;
     PyObject *py_unicode = NULL;
     LPWSTR *szArglist = NULL;
     int nArgs, i;
 
-    if (psutil_get_process_data(pid, KIND_CMDLINE, &data) != 0)
+    if (psutil_get_process_data(pid, KIND_CMDLINE, &data, &size) != 0)
         goto out;
 
     // attempt to parse the command line using Win32 API
@@ -622,8 +731,9 @@ out:
 PyObject *psutil_get_cwd(long pid) {
     PyObject *ret = NULL;
     WCHAR *data = NULL;
+    SIZE_T size;
 
-    if (psutil_get_process_data(pid, KIND_CWD, &data) != 0)
+    if (psutil_get_process_data(pid, KIND_CWD, &data, &size) != 0)
         goto out;
 
     // convert wchar array to a Python unicode string
@@ -635,6 +745,26 @@ out:
     return ret;
 }
 
+/*
+ * returns a Python string containing the environment variable data for the
+ * process with given pid or NULL on error.
+ */
+PyObject *psutil_get_environ(long pid) {
+    PyObject *ret = NULL;
+    WCHAR *data = NULL;
+    SIZE_T size;
+
+    if (psutil_get_process_data(pid, KIND_ENVIRON, &data, &size) != 0)
+        goto out;
+
+    // convert wchar array to a Python unicode string
+    ret = PyUnicode_FromWideChar(data, size / 2);
+
+out:
+    free(data);
+
+    return ret;
+}
 
 #define PH_FIRST_PROCESS(Processes) ((PSYSTEM_PROCESS_INFORMATION)(Processes))
 #define PH_NEXT_PROCESS(Process) ( \
diff --git a/psutil/arch/windows/process_info.h b/psutil/arch/windows/process_info.h
index 908efc5f..f9af7765 100644
--- a/psutil/arch/windows/process_info.h
+++ b/psutil/arch/windows/process_info.h
@@ -20,6 +20,7 @@ int psutil_pid_in_proclist(DWORD pid);
 int psutil_pid_is_running(DWORD pid);
 PyObject* psutil_get_cmdline(long pid);
 PyObject* psutil_get_cwd(long pid);
+PyObject* psutil_get_environ(long pid);
 int psutil_get_proc_info(DWORD pid, PSYSTEM_PROCESS_INFORMATION *retProcess,
                          PVOID *retBuffer);
 
diff --git a/test/_windows.py b/test/_windows.py
index f812880c..879a698c 100644
--- a/test/_windows.py
+++ b/test/_windows.py
@@ -526,6 +526,8 @@ class RemoteProcessTestCase(unittest.TestCase):
     test_args = ["-c", "import sys; sys.stdin.read()"]
 
     def setUp(self):
+        env = os.environ.copy()
+        env["THINK_OF_A_NUMBER"] = str(os.getpid())
         self.proc32 = get_test_subprocess([self.python32] + self.test_args,
                                           env=env,
                                           stdin=subprocess.PIPE)
@@ -560,6 +562,18 @@ class RemoteProcessTestCase(unittest.TestCase):
         p = psutil.Process(self.proc64.pid)
         self.assertEqual(p.cwd(), os.getcwd())
 
+    def test_environ_32(self):
+        p = psutil.Process(self.proc32.pid)
+        e = p.environ()
+        self.assertIn("THINK_OF_A_NUMBER", e)
+        self.assertEquals(e["THINK_OF_A_NUMBER"], str(os.getpid()))
+
+    def test_environ_64(self):
+        p = psutil.Process(self.proc64.pid)
+        e = p.environ()
+        self.assertIn("THINK_OF_A_NUMBER", e)
+        self.assertEquals(e["THINK_OF_A_NUMBER"], str(os.getpid()))
+
 
 def main():
     test_suite = unittest.TestSuite()
diff --git a/test/test_memory_leaks.py b/test/test_memory_leaks.py
index a43e941e..d79d5900 100644
--- a/test/test_memory_leaks.py
+++ b/test/test_memory_leaks.py
@@ -323,7 +323,7 @@ class TestProcessObjectLeaks(Base):
                 s.close()
 
     @unittest.skipUnless(hasattr(psutil.Process, 'environ'),
-                         "Linux and OSX")
+                         "Linux, OSX and Windows")
     def test_environ(self):
         self.execute("environ")