#1398 #1348 / win / cmdline: refactor code so that the 2 cmdline() im…

…plementations can be called separately (and tested separately)
giampaolo · Feb 15, 2019 · f1374c3 · f1374c3
1 parent 3f9643f
commit f1374c3
Show file tree

Hide file tree

Showing 5 changed files with 51 additions and 25 deletions.
diff --git a/psutil/_psutil_windows.c b/psutil/_psutil_windows.c
@@ -708,22 +708,28 @@ psutil_cpu_count_phys(PyObject *self, PyObject *args) {
  * 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
@@ -708,7 +708,13 @@ def exe(self):
 
     @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
@@ -931,7 +931,7 @@ psutil_get_cmdline_data(long pid, WCHAR **pdata, SIZE_T *psize) {
  * 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
@@ -384,6 +384,16 @@ def test_proc_info(self):
         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
@@ -673,6 +673,22 @@ def test_num_handles(self):
                              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):