From 6fe1544a7e2054c2eb16bf750b4a764e7fa87718 Mon Sep 17 00:00:00 2001 From: Frank Benkstein Date: Thu, 28 Jan 2016 21:59:12 +0100 Subject: implement Process.environ() on Windows --- psutil/__init__.py | 2 +- psutil/_psutil_windows.c | 25 +++++++ psutil/_psutil_windows.h | 1 + psutil/_pswindows.py | 5 ++ psutil/arch/windows/process_info.c | 140 +++++++++++++++++++++++++++++++++++-- psutil/arch/windows/process_info.h | 1 + test/_windows.py | 14 ++++ test/test_memory_leaks.py | 2 +- 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 @@ -591,6 +591,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. */ @@ -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") -- cgit v1.2.1