Repository: giampaolo/psutil
Pull request #732 (2016-01-19 19:05:45)

Revive Process.environ
Hi,

here's another shot at implementing Process.environ.  This has been running productively at my employer for a couple of months which shook out a few bugs.  We're using this to track subprocess trees and terminate a whole hierarchy by marking them with an environment variable.  This is similar to Jenkins' process tree killer.  Please have a look.

Best regards,
Frank.

fbenkstein: This looks to be ready for review.
giampaolo: This is cool, thanks for working on this. FYI in case you don't know there's some previous work that had been done here: https://code.google.com/p/psutil/issues/detail?id=52
Also, some confusing code regarding Solaris:
https://code.google.com/p/psutil/issues/detail?id=247

One question for you: in here (https://code.google.com/p/psutil/issues/detail?id=52#c18) I stated that on Windows you can get environ only for the current process (or else AccessDenied exception). Does this patch fix that issue or not?
fbenkstein: I saw the earlier tries and used it early in my implementation.  Because of the problems on Windows I took the code from the blog post mentioned in the commit message instead.  I can confirm that the Windows code works not just for the current process.  I think there was a wrong offset or address in the earlier implementation when calling ReadProcessMemory.

Should I put this into separate PRs so it can be reviewed separately?
giampaolo: > I can confirm that the Windows code works not just for the current process. I think there was 
> a wrong offset or address in the earlier implementation when calling ReadProcessMemory.

Great! This was one of the main reasons which led me to give up on this.

> Should I put this into separate PRs so it can be reviewed separately?

I don't think it's necessary. 
fbenkstein: I'm still not sure how to handle non-utf8 data consistently between Python 2 and 3.  Should it be byte strings in both cases?  I don't have much experience with Python 3 but it seems that users there seem to expect str (i.e. unicode) and the rest of the psutil seems to be modelled that way.  Unfortunately that means that process names, command lines and now environments that contain non-utf8 data will break applications.

Example:
    cp /bin/cat $'\xff\xfe'
    pid=$!
    python2.7 -c "import psutil; print(psutil.Process($pid).name())"
    ▒▒
    python3.5 -c "import psutil; print(psutil.Process($pid).name())"
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/xxx/Projects/psutil/psutil/__init__.py", line 543, in name
    name = self._proc.name()
  File "/home/xxx/Projects/psutil/psutil/_pslinux.py", line 811, in wrapper
    return fun(self, *args, **kwargs)
  File "/home/xxx/Projects/psutil/psutil/_pslinux.py", line 855, in name
    data = f.read()
  File "/usr/lib/python3.5/codecs.py", line 321, in decode
    (result, consumed) = self._buffer_decode(data, self.errors, final)
UnicodeDecodeError: 'utf-8' codec can't decode byte 0xff in position 7: invalid start byte


fbenkstein: I've side stepped the encoding issue for now.  Also, I think it's better if I move the Windows code into a separate PR if that's okay with you.
giampaolo: Sure, make another PR if you want.
fbenkstein: Okay.  I'm done for today.  Will try to prepare the Windows change tomorrow.
fbenkstein: The AppVeyor test failures seem unrelated.  Is there anything else missing here?
giampaolo: I added a couple more comments, the rest is good. 
giampaolo: OSX failure https://travis-ci.org/giampaolo/psutil/jobs/104056877:

```
'/private/var/folders/gw/_2jq29095y7b__wtby9dg_5h0000gn/T/tmpnGpanB/\xc0\x80'
```
vs:
```
             '/var/folders/gw/_2jq29095y7b__wtby9dg_5h0000gn/T/tmpnGpanB/\xc0\x80'
```
I currently don't have access to OSX but I seem to remember this happens because /private is a symlink so perhaps you should be able to fix this by using `os.path.realpath` which resolves symlinks.

Patch:
diff --git a/docs/index.rst b/docs/index.rst
index bff8d6ed..566b9996 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -693,6 +693,15 @@ Process class
 
      The command line this process has been called with.
 
+  .. method:: environ()
+
+     The environment variables of the process as a dict.  Note: this might not
+     reflect changes made after the process started.
+
+     Availability: Linux, OSX
+
+     .. versionchanged: 3.4.3: added
+
   .. method:: create_time()
 
      The process creation time as a floating point number expressed in seconds
diff --git a/psutil/__init__.py b/psutil/__init__.py
index 10eee585..c034ad6f 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -744,6 +744,13 @@ class Process(object):
             else:
                 self._proc.cpu_affinity_set(list(set(cpus)))
 
+    # Linux and OSX only
+    if hasattr(_psplatform.Process, "environ"):
+        def environ(self):
+            """The environment variables of the process as a dict.  Note: this
+            might not reflect changes made after the process started.  """
+            return self._proc.environ()
+
     if _WINDOWS:
 
         def num_handles(self):
diff --git a/psutil/_common.py b/psutil/_common.py
index 3ff1cb88..d8c30190 100644
--- a/psutil/_common.py
+++ b/psutil/_common.py
@@ -169,6 +169,32 @@ def supports_ipv6():
         return False
 
 
+def parse_environ_block(data):
+    """Parse a C environ block of environment variables into a dictionary."""
+    # The block is usually raw data from the target process.  It might contain
+    # trailing garbage and lines that do not look like assignments.
+    ret = {}
+    pos = 0
+
+    while True:
+        next_pos = data.find("\0", pos)
+        # nul byte at the beginning or double nul byte means finish
+        if next_pos <= pos:
+            break
+        # there might not be an equals sign
+        equal_pos = data.find("=", pos, next_pos)
+        if equal_pos > pos:
+            key = data[pos:equal_pos]
+            value = data[equal_pos+1:next_pos]
+            # Windows expects environment variables to be uppercase only
+            if os.name == "nt":
+                key = key.upper()
+            ret[key] = value
+        pos = next_pos + 1
+
+    return ret
+
+
 def sockfam_to_enum(num):
     """Convert a numeric socket family value to an IntEnum member.
     If it's not a known member, return the numeric value itself.
diff --git a/psutil/_pslinux.py b/psutil/_pslinux.py
index 2f06e1e9..ee613564 100644
--- a/psutil/_pslinux.py
+++ b/psutil/_pslinux.py
@@ -25,6 +25,7 @@ from . import _psutil_linux as cext
 from . import _psutil_posix as cext_posix
 from ._common import isfile_strict
 from ._common import memoize
+from ._common import parse_environ_block
 from ._common import NIC_DUPLEX_FULL
 from ._common import NIC_DUPLEX_HALF
 from ._common import NIC_DUPLEX_UNKNOWN
@@ -886,6 +887,12 @@ class Process(object):
             data = data[:-1]
         return [x for x in data.split('\x00')]
 
+    @wrap_exceptions
+    def environ(self):
+        with open_text("%s/%s/environ" % (self._procfs_path, self.pid)) as f:
+            data = f.read()
+        return parse_environ_block(data)
+
     @wrap_exceptions
     def terminal(self):
         tmap = _psposix._get_terminal_map()
diff --git a/psutil/_psosx.py b/psutil/_psosx.py
index 0e770b80..271e31f9 100644
--- a/psutil/_psosx.py
+++ b/psutil/_psosx.py
@@ -15,6 +15,7 @@ from . import _psutil_osx as cext
 from . import _psutil_posix as cext_posix
 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
@@ -237,6 +238,12 @@ class Process(object):
             raise NoSuchProcess(self.pid, self._name)
         return cext.proc_cmdline(self.pid)
 
+    @wrap_exceptions
+    def environ(self):
+        if not pid_exists(self.pid):
+            raise NoSuchProcess(self.pid, self._name)
+        return parse_environ_block(cext.proc_environ(self.pid))
+
     @wrap_exceptions
     def ppid(self):
         return cext.proc_ppid(self.pid)
diff --git a/psutil/_psutil_osx.c b/psutil/_psutil_osx.c
index 7f8c0ff7..4ffae898 100644
--- a/psutil/_psutil_osx.c
+++ b/psutil/_psutil_osx.c
@@ -214,6 +214,23 @@ psutil_proc_cmdline(PyObject *self, PyObject *args) {
 }
 
 
+/*
+ * Return process environment as a Python string.
+ */
+static PyObject *
+psutil_proc_environ(PyObject *self, PyObject *args) {
+    long pid;
+    PyObject *py_retdict = NULL;
+
+    if (! PyArg_ParseTuple(args, "l", &pid))
+        return NULL;
+
+    // get the environment block, defined in arch/osx/process_info.c
+    py_retdict = psutil_get_environ(pid);
+    return py_retdict;
+}
+
+
 /*
  * Return process parent pid from kinfo_proc as a Python integer.
  */
@@ -1664,6 +1681,8 @@ PsutilMethods[] = {
      "Return process name"},
     {"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_cwd", psutil_proc_cwd, METH_VARARGS,
diff --git a/psutil/_psutil_osx.h b/psutil/_psutil_osx.h
index 907a8e53..2220ec3b 100644
--- a/psutil/_psutil_osx.h
+++ b/psutil/_psutil_osx.h
@@ -12,6 +12,7 @@ static PyObject* psutil_proc_connections(PyObject* self, PyObject* args);
 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_environ(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_exe(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_gids(PyObject* self, PyObject* args);
 static PyObject* psutil_proc_memory_info(PyObject* self, PyObject* args);
diff --git a/psutil/arch/osx/process_info.c b/psutil/arch/osx/process_info.c
index 20aafdac..4c4ec88d 100644
--- a/psutil/arch/osx/process_info.c
+++ b/psutil/arch/osx/process_info.c
@@ -238,6 +238,124 @@ error:
 }
 
 
+// return process environment as a python string
+PyObject *
+psutil_get_environ(long pid) {
+    int mib[3];
+    int nargs;
+    char *procargs = NULL;
+    char *procenv = NULL;
+    char *arg_ptr;
+    char *arg_end;
+    char *env_start;
+    size_t argmax;
+    PyObject *py_ret = NULL;
+
+    // special case for PID 0 (kernel_task) where cmdline cannot be fetched
+    if (pid == 0)
+        goto empty;
+
+    // read argmax and allocate memory for argument space.
+    argmax = psutil_get_argmax();
+    if (! argmax) {
+        PyErr_SetFromErrno(PyExc_OSError);
+        goto error;
+    }
+
+    procargs = (char *)malloc(argmax);
+    if (NULL == procargs) {
+        PyErr_SetFromErrno(PyExc_OSError);
+        goto error;
+    }
+
+    // read argument space
+    mib[0] = CTL_KERN;
+    mib[1] = KERN_PROCARGS2;
+    mib[2] = pid;
+    if (sysctl(mib, 3, procargs, &argmax, NULL, 0) < 0) {
+        if (EINVAL == errno) {
+            // EINVAL == access denied OR nonexistent PID
+            if (psutil_pid_exists(pid))
+                AccessDenied();
+            else
+                NoSuchProcess();
+        }
+        goto error;
+    }
+
+    arg_end = &procargs[argmax];
+    // copy the number of arguments to nargs
+    memcpy(&nargs, procargs, sizeof(nargs));
+
+    // skip executable path
+    arg_ptr = procargs + sizeof(nargs);
+    arg_ptr = memchr(arg_ptr, '\0', arg_end - arg_ptr);
+
+    if (arg_ptr == NULL || arg_ptr == arg_end)
+        goto empty;
+
+    // skip ahead to the first argument
+    for (; arg_ptr < arg_end; arg_ptr++) {
+        if (*arg_ptr != '\0')
+            break;
+    }
+
+    // iterate through arguments
+    while (arg_ptr < arg_end && nargs > 0) {
+        if (*arg_ptr++ == '\0')
+            nargs--;
+    }
+
+    // build an environment variable block
+    env_start = arg_ptr;
+
+    procenv = calloc(1, arg_end - arg_ptr);
+
+    if (procenv == NULL) {
+        PyErr_NoMemory();
+        goto error;
+    }
+
+    while (*arg_ptr != '\0' && arg_ptr < arg_end) {
+        char *s = memchr(arg_ptr + 1, '\0', arg_end - arg_ptr);
+
+        if (s == NULL)
+            break;
+
+        memcpy(procenv + (arg_ptr - env_start), arg_ptr, s - arg_ptr);
+
+        arg_ptr = s + 1;
+    }
+
+#if PY_MAJOR_VERSION >= 3
+    py_ret = PyUnicode_FromStringAndSize(procenv, arg_ptr - env_start + 1);
+#else
+    py_ret = PyString_FromStringAndSize(procenv, arg_ptr - env_start + 1);
+#endif
+
+    if (!py_ret)
+        goto error;
+
+    free(procargs);
+    free(procenv);
+
+    return py_ret;
+
+empty:
+    if (procargs != NULL)
+        free(procargs);
+    return Py_BuildValue("s", "");
+
+error:
+    Py_XDECREF(py_ret);
+    if (procargs != NULL)
+        free(procargs);
+    if (procenv != NULL)
+        free(procargs);
+    return NULL;
+}
+
+
 int
 psutil_get_kinfo_proc(pid_t pid, struct kinfo_proc *kp) {
     int mib[4];
diff --git a/psutil/arch/osx/process_info.h b/psutil/arch/osx/process_info.h
index ac71bbe3..8bc10ec1 100644
--- a/psutil/arch/osx/process_info.h
+++ b/psutil/arch/osx/process_info.h
@@ -14,3 +14,4 @@ int psutil_get_proc_list(kinfo_proc **procList, size_t *procCount);
 int psutil_pid_exists(long pid);
 int psutil_proc_pidinfo(long pid, int flavor, void *pti, int size);
 PyObject* psutil_get_cmdline(long pid);
+PyObject* psutil_get_environ(long pid);


Test patch:
diff --git a/test/test_memory_leaks.py b/test/test_memory_leaks.py
index d971513d..a43e941e 100644
--- a/test/test_memory_leaks.py
+++ b/test/test_memory_leaks.py
@@ -322,6 +322,11 @@ class TestProcessObjectLeaks(Base):
             for s in socks:
                 s.close()
 
+    @unittest.skipUnless(hasattr(psutil.Process, 'environ'),
+                         "Linux and OSX")
+    def test_environ(self):
+        self.execute("environ")
+
 
 p = get_test_subprocess()
 DEAD_PROC = psutil.Process(p.pid)
diff --git a/test/test_psutil.py b/test/test_psutil.py
index c72a8543..d5f52e56 100644
--- a/test/test_psutil.py
+++ b/test/test_psutil.py
@@ -507,7 +507,7 @@ def skip_on_not_implemented(only_if=None):
     return decorator
 
 
-def create_temp_executable_file(suffix):
+def create_temp_executable_file(suffix, code="void main() { pause(); }"):
     tmpdir = None
     if TRAVIS and OSX:
         tmpdir = "/private/tmp"
@@ -520,7 +520,7 @@ def create_temp_executable_file(suffix):
             prefix='psu', suffix='.c', dir=tmpdir)
         os.close(fd)
         with open(c_file, "w") as f:
-            f.write("void main() { pause(); }")
+            f.write(code)
         subprocess.check_call(["gcc", c_file, "-o", path])
         safe_remove(c_file)
     else:
@@ -2500,6 +2500,57 @@ class TestProcess(unittest.TestCase):
             proc.wait()
             self.assertIsNotNone(proc.returncode)
 
+    @unittest.skipUnless(hasattr(psutil.Process, "environ"),
+                         "environ not available")
+    def test_environ(self):
+        self.maxDiff = None
+        p = psutil.Process()
+        d = p.environ()
+        d2 = os.environ.copy()
+
+        if OSX:
+            for key in (
+                "__CF_USER_TEXT_ENCODING",
+                "VERSIONER_PYTHON_PREFER_32_BIT",
+                "VERSIONER_PYTHON_VERSION",
+            ):
+                d.pop(key, None)
+                d2.pop(key, None)
+
+        self.assertEqual(d, d2)
+
+    @unittest.skipUnless(hasattr(psutil.Process, "environ"),
+                         "environ not available")
+    @unittest.skipUnless(POSIX, "posix only")
+    def test_weird_environ(self):
+        # environment variables can contain values without an equals sign
+        code = textwrap.dedent("""
+        #include <unistd.h>
+        #include <fcntl.h>
+        char * const argv[] = {"cat", 0};
+        char * const envp[] = {"A=1", "X", "C=3", 0};
+        int main(void) {
+            /* Close stderr on exec so parent can wait for the execve to
+             * finish. */
+            if (fcntl(2, F_SETFD, FD_CLOEXEC) != 0)
+                return 0;
+            return execve("/bin/cat", argv, envp);
+        }
+        """)
+        path = create_temp_executable_file("x", code=code)
+        self.addCleanup(safe_remove, path)
+        sproc = get_test_subprocess([path],
+                                    stdin=subprocess.PIPE,
+                                    stderr=subprocess.PIPE)
+        p = psutil.Process(sproc.pid)
+        wait_for_pid(p.pid)
+        self.assertTrue(p.is_running())
+        # Wait for process to exec or exit.
+        self.assertEqual(sproc.stderr.read(), b"")
+        self.assertEqual(p.environ(), {"A": "1", "C": "3"})
+        sproc.communicate()
+        self.assertEqual(sproc.returncode, 0)
+
 
 # ===================================================================
 # --- Featch all processes test
@@ -2561,7 +2612,7 @@ class TestFetchAllProcesses(unittest.TestCase):
                         self.assertTrue(str(err))
                         self.assertTrue(err.msg)
                     else:
-                        if ret not in (0, 0.0, [], None, ''):
+                        if ret not in (0, 0.0, [], None, '', {}):
                             assert ret, ret
                         meth = getattr(self, name)
                         meth(ret, p)
@@ -2774,6 +2825,9 @@ class TestFetchAllProcesses(unittest.TestCase):
         self.assertGreaterEqual(ret[0], -1)
         self.assertGreaterEqual(ret[1], -1)
 
+    def environ(self, ret, proc):
+        self.assertIsInstance(ret, dict)
+
 
 # ===================================================================
 # --- Limited user tests
@@ -2999,6 +3053,27 @@ class TestMisc(unittest.TestCase):
         # docstring
         self.assertEqual(foo.__doc__, "foo docstring")
 
+    def test_parse_environ_block(self):
+        from psutil._common import parse_environ_block
+
+        def k(s):
+            return s.upper() if WINDOWS else s
+
+        self.assertEqual(parse_environ_block("a=1\0"),
+                         {k("a"): "1"})
+        self.assertEqual(parse_environ_block("a=1\0b=2\0\0"),
+                         {k("a"): "1", k("b"): "2"})
+        self.assertEqual(parse_environ_block("a=1\0b=\0\0"),
+                         {k("a"): "1", k("b"): ""})
+        # ignore everything after \0\0
+        self.assertEqual(parse_environ_block("a=1\0b=2\0\0c=3\0"),
+                         {k("a"): "1", k("b"): "2"})
+        # ignore everything that is not an assignment
+        self.assertEqual(parse_environ_block("xxx\0a=1\0"), {k("a"): "1"})
+        self.assertEqual(parse_environ_block("a=1\0=b=2\0"), {k("a"): "1"})
+        # do not fail if the block is incomplete
+        self.assertEqual(parse_environ_block("a=1\0b=2"), {k("a"): "1"})
+
     def test_supports_ipv6(self):
         if supports_ipv6():
             with mock.patch('psutil._common.socket') as s: