Fix #1494: [AIX] implement Process.environ() (#1505) (patch by Arnon Yaari)

4 files changed, 80 insertions, 2 deletions

diff --git a/psutil/__init__.py b/psutil/__init__.py
index bd968f5d..5d37c5c0 100644
--- a/psutil/__init__.py
+++ b/psutil/__init__.py
@@ -949,7 +949,7 @@ class Process(object):
             """
             return self._proc.cpu_num()
 
-    # Linux, macOS and Windows only
+    # Linux, macOS, Windows, Solaris, AIX
     if hasattr(_psplatform.Process, "environ"):
 
         def environ(self):
diff --git a/psutil/_psaix.py b/psutil/_psaix.py
index d0aa0e94..5810e3b4 100644
--- a/psutil/_psaix.py
+++ b/psutil/_psaix.py
@@ -414,6 +414,10 @@ class Process(object):
         return cext.proc_args(self.pid)
 
     @wrap_exceptions
+    def environ(self):
+        return cext.proc_environ(self.pid)
+
+    @wrap_exceptions
     def create_time(self):
         return self._proc_basic_info()[proc_info_map['create_time']]
 
diff --git a/psutil/_psutil_aix.c b/psutil/_psutil_aix.c
index 1a83ac24..723d159d 100644
--- a/psutil/_psutil_aix.c
+++ b/psutil/_psutil_aix.c
@@ -213,6 +213,78 @@ error:
 }
 
 
+/*
+ * Return process environment variables as a Python dict
+ */
+static PyObject *
+psutil_proc_environ(PyObject *self, PyObject *args) {
+    int pid;
+    PyObject *py_retdict = PyDict_New();
+    PyObject *py_key = NULL;
+    PyObject *py_val = NULL;
+    struct procsinfo procbuf;
+    long env_max;
+    char *envbuf = NULL;
+    char *curvar = NULL;
+    char *separator = NULL;
+    int ret;
+
+    if (py_retdict == NULL)
+        return NULL;
+    if (!PyArg_ParseTuple(args, "i", &pid))
+        goto error;
+    env_max = sysconf(_SC_ARG_MAX);
+    envbuf = malloc(env_max);
+    if (envbuf == NULL) {
+        PyErr_NoMemory();
+        goto error;
+    }
+
+    procbuf.pi_pid = pid;
+    ret = getevars(&procbuf, sizeof(struct procinfo), envbuf, ARG_MAX);
+    if (ret == -1) {
+        PyErr_SetFromErrno(PyExc_OSError);
+        goto error;
+    }
+
+    curvar = envbuf;
+    /* getevars will always append an extra NULL to end the arg list,
+     * even if the buffer is not big enough (even though it is supposed
+     * to be) so the following 'while' is safe */
+    while (*curvar != '\0') {
+        separator = strchr(curvar, '=');
+        if (separator != NULL) {
+            py_key = PyUnicode_DecodeFSDefaultAndSize(
+                curvar,
+                (Py_ssize_t)(separator - curvar)
+            );
+            if (!py_key)
+                goto error;
+            py_val = PyUnicode_DecodeFSDefault(separator + 1);
+            if (!py_val)
+                goto error;
+            if (PyDict_SetItem(py_retdict, py_key, py_val))
+                goto error;
+            Py_DECREF(py_key);
+            Py_DECREF(py_val);
+        }
+        curvar = strchr(curvar, '\0') + 1;
+    }
+
+    free(envbuf);
+
+    return py_retdict;
+
+error:
+    if (envbuf != NULL)
+        free(envbuf);
+    Py_XDECREF(py_retdict);
+    Py_XDECREF(py_key);
+    Py_XDECREF(py_val);
+    return NULL;
+}
+
+
 #ifdef CURR_VERSION_THREAD
 
 /*
@@ -924,6 +996,8 @@ PsutilMethods[] =
      "Return process name."},
     {"proc_args", psutil_proc_args, METH_VARARGS,
      "Return process command line arguments."},
+    {"proc_environ", psutil_proc_environ, METH_VARARGS,
+     "Return process environment variables."},
     {"proc_cpu_times", psutil_proc_cpu_times, METH_VARARGS,
      "Return process user and system CPU times."},
     {"proc_cred", psutil_proc_cred, METH_VARARGS,
diff --git a/psutil/tests/test_contracts.py b/psutil/tests/test_contracts.py
index adf7b680..20da8241 100755
--- a/psutil/tests/test_contracts.py
+++ b/psutil/tests/test_contracts.py
@@ -130,7 +130,7 @@ class TestAvailability(unittest.TestCase):
 
     def test_proc_environ(self):
         self.assertEqual(hasattr(psutil.Process, "environ"),
-                         LINUX or MACOS or WINDOWS)
+                         LINUX or MACOS or WINDOWS or AIX or SUNOS)
 
     def test_proc_uids(self):
         self.assertEqual(hasattr(psutil.Process, "uids"), POSIX)