revive Process.environ on OSX

Revive Process.environ on OSX.  The data is right behind the command line so
the new function psutil_get_environ is mostly a copy of psutil_get_cmdline.

8 files changed, 161 insertions, 3 deletions

diff --git a/psutil/__init__.py b/psutil/__init__.py
index cec43c75..c034ad6f 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 only
+    # Linux and OSX only
     if hasattr(_psplatform.Process, "environ"):
         def environ(self):
             """The environment variables of the process as a dict.  Note: this
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
@@ -238,6 +239,12 @@ class Process(object):
         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
@@ -215,6 +215,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.
  */
 static PyObject *
@@ -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);
diff --git a/test/test_memory_leaks.py b/test/test_memory_leaks.py
index b5112fa2..a43e941e 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 only")
+                         "Linux and OSX")
     def test_environ(self):
         self.execute("environ")
 
diff --git a/test/test_psutil.py b/test/test_psutil.py
index 55be9714..d5f52e56 100644
--- a/test/test_psutil.py
+++ b/test/test_psutil.py
@@ -2505,7 +2505,19 @@ class TestProcess(unittest.TestCase):
     def test_environ(self):
         self.maxDiff = None
         p = psutil.Process()
-        self.assertEqual(p.environ(), os.environ)
+        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")