diff options
author | Giampaolo Rodola <g.rodola@gmail.com> | 2019-02-15 10:30:08 -0800 |
---|---|---|
committer | Giampaolo Rodola <g.rodola@gmail.com> | 2019-02-15 10:30:08 -0800 |
commit | f1374c36d07abd3aaa15d6f512143189c1748099 (patch) | |
tree | ac10909e7d892618a1a3a4802e06363c6f808b6a | |
parent | 3f9643f2357e887bff3ab4ecf093b9f28fa53fe9 (diff) | |
download | psutil-f1374c36d07abd3aaa15d6f512143189c1748099.tar.gz |
#1398 #1348 / win / cmdline: refactor code so that the 2 cmdline() implementations can be called separately (and tested separately)
-rw-r--r-- | psutil/_psutil_windows.c | 15 | ||||
-rw-r--r-- | psutil/_pswindows.py | 8 | ||||
-rw-r--r-- | psutil/arch/windows/process_info.c | 27 | ||||
-rwxr-xr-x | psutil/tests/test_memory_leaks.py | 10 | ||||
-rwxr-xr-x | psutil/tests/test_windows.py | 16 |
5 files changed, 51 insertions, 25 deletions
diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c index 4251e0c7..eb35c5f7 100644 --- a/psutil/_psutil_windows.c +++ b/psutil/_psutil_windows.c @@ -708,22 +708,28 @@ return_none: * Return process cmdline as a Python list of cmdline arguments. */ static PyObject * -psutil_proc_cmdline(PyObject *self, PyObject *args) { +psutil_proc_cmdline(PyObject *self, PyObject *args, PyObject *kwdict) { long pid; int pid_return; + int use_peb; + PyObject *py_usepeb = Py_True; + static char *keywords[] = {"pid", "use_peb", NULL}; - if (! PyArg_ParseTuple(args, "l", &pid)) + if (!PyArg_ParseTupleAndKeywords(args, kwdict, "i|O", + keywords, &pid, &py_usepeb)) { return NULL; + } if ((pid == 0) || (pid == 4)) return Py_BuildValue("[]"); + use_peb = (py_usepeb == Py_True); pid_return = psutil_pid_is_running(pid); if (pid_return == 0) return NoSuchProcess(""); if (pid_return == -1) return NULL; - return psutil_get_cmdline(pid); + return psutil_get_cmdline(pid, use_peb); } @@ -3688,8 +3694,7 @@ psutil_sensors_battery(PyObject *self, PyObject *args) { static PyMethodDef PsutilMethods[] = { // --- per-process functions - - {"proc_cmdline", psutil_proc_cmdline, METH_VARARGS, + {"proc_cmdline", (PyCFunction)(void(*)(void))psutil_proc_cmdline, METH_VARARGS | METH_KEYWORDS, "Return process cmdline as a list of cmdline arguments"}, {"proc_environ", psutil_proc_environ, METH_VARARGS, "Return process environment data"}, diff --git a/psutil/_pswindows.py b/psutil/_pswindows.py index 1aeb46ef..e66febe0 100644 --- a/psutil/_pswindows.py +++ b/psutil/_pswindows.py @@ -708,7 +708,13 @@ class Process(object): @wrap_exceptions def cmdline(self): - ret = cext.proc_cmdline(self.pid) + try: + ret = cext.proc_cmdline(self.pid, use_peb=True) + except OSError as err: + if err.errno in ACCESS_DENIED_ERRSET: + ret = cext.proc_cmdline(self.pid, use_peb=False) + else: + raise if PY3: return ret else: diff --git a/psutil/arch/windows/process_info.c b/psutil/arch/windows/process_info.c index 051b1c6d..6f877342 100644 --- a/psutil/arch/windows/process_info.c +++ b/psutil/arch/windows/process_info.c @@ -931,7 +931,7 @@ error: * with given pid or NULL on error. */ PyObject * -psutil_get_cmdline(long pid) { +psutil_get_cmdline(long pid, int use_peb) { PyObject *ret = NULL; WCHAR *data = NULL; SIZE_T size; @@ -956,25 +956,14 @@ psutil_get_cmdline(long pid) { - https://github.com/giampaolo/psutil/pull/1398 - https://blog.xpnsec.com/how-to-argue-like-cobalt-strike/ */ - func_ret = psutil_get_process_data(pid, KIND_CMDLINE, &data, &size); - if (func_ret != 0) { - if ((GetLastError() == ERROR_ACCESS_DENIED) && - (windows_version >= WINDOWS_81)) - { - // reset that we had an error - // and retry with NtQueryInformationProcess - // (for protected processes) - PyErr_Clear(); - - func_ret = psutil_get_cmdline_data(pid, &data, &size); - if (func_ret != 0) { - goto out; - } - } - else { - goto out; - } + if (use_peb == 1) { + func_ret = psutil_get_process_data(pid, KIND_CMDLINE, &data, &size); + } + else { + func_ret = psutil_get_cmdline_data(pid, &data, &size); } + if (func_ret != 0) + goto out; // attempt to parse the command line using Win32 API szArglist = CommandLineToArgvW(data, &nArgs); diff --git a/psutil/tests/test_memory_leaks.py b/psutil/tests/test_memory_leaks.py index fc3a0365..58ae3312 100755 --- a/psutil/tests/test_memory_leaks.py +++ b/psutil/tests/test_memory_leaks.py @@ -384,6 +384,16 @@ class TestProcessObjectLeaks(TestMemLeak): self.execute(cext.proc_info, os.getpid()) +class TestProcessDualImplementation(TestMemLeak): + + if WINDOWS: + def test_cmdline_peb_true(self): + self.execute(cext.proc_cmdline, os.getpid(), use_peb=True) + + def test_cmdline_peb_false(self): + self.execute(cext.proc_cmdline, os.getpid(), use_peb=False) + + class TestTerminatedProcessLeaks(TestProcessObjectLeaks): """Repeat the tests above looking for leaks occurring when dealing with terminated processes raising NoSuchProcess exception. diff --git a/psutil/tests/test_windows.py b/psutil/tests/test_windows.py index 4633f759..35ea6217 100755 --- a/psutil/tests/test_windows.py +++ b/psutil/tests/test_windows.py @@ -673,6 +673,22 @@ class TestDualProcessImplementation(unittest.TestCase): num_handles) assert fun.called + def test_cmdline(self): + from psutil._pswindows import ACCESS_DENIED_ERRSET + for pid in psutil.pids(): + try: + a = cext.proc_cmdline(pid, use_peb=True) + b = cext.proc_cmdline(pid, use_peb=False) + except OSError as err: + if err.errno in ACCESS_DENIED_ERRSET: + pass + elif err.errno == errno.ESRCH: + pass # NSP + else: + raise + else: + self.assertEqual(a, b) + @unittest.skipIf(not WINDOWS, "WINDOWS only") class RemoteProcessTestCase(unittest.TestCase): |