revive Process.environ on Linux

Revive Process.environ on Linux.  Implemented by reading from
/proc/<pid>/environ.  Parsing is put into a helper function
parse_environ_block to be able to reuse it on other platforms.

5 files changed, 111 insertions, 3 deletions

diff --git a/psutil/__init__.py b/psutil/__init__.py
index 10eee585..cec43c75 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 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
@@ -887,6 +888,12 @@ class Process(object):
         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()
         with open_binary("%s/%s/stat" % (self._procfs_path, self.pid)) as f:
diff --git a/test/test_memory_leaks.py b/test/test_memory_leaks.py
index d971513d..b5112fa2 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 only")
+    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..55be9714 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,45 @@ 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()
+        self.assertEqual(p.environ(), os.environ)
+
+    @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 +2600,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 +2813,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 +3041,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: