[tor-commits] [stem/master] Add function pids_by_user()

[atagar at torproject.org]

commit ae178869ee37a1ec78e6b7ef511b54ca4f74863e
Author: Damian Johnson <atagar at torproject.org>
Date:   Sat Jan 30 14:35:52 2016 -0800

    Add function pids_by_user()
    
    Simple utility for the system module to get pids for a user.
---
 docs/change_log.rst       |  1 +
 stem/util/system.py       | 31 +++++++++++++++++++++++++++++++
 test/integ/util/system.py | 10 ++++++++++
 3 files changed, 42 insertions(+)

diff --git a/docs/change_log.rst b/docs/change_log.rst
index 74159e2..81bf5a2 100644
--- a/docs/change_log.rst
+++ b/docs/change_log.rst
@@ -77,6 +77,7 @@ The following are only available within Stem's `git repository
   * Allow :func:`stem.util.conf.Config.set` to remove values when provided with a **None** value
   * Additional information when :func:`~stem.util.system.call` fails through a :class:`~stem.util.system.CallError`
   * Added an **is_ipv6** value to :class:`~stem.util.connection.Connection` instances
+  * Added :func:`~stem.util.system.pids_by_user`
   * Added :func:`~stem.util.__init__.datetime_to_unix`
 
  * **Interpreter**
diff --git a/stem/util/system.py b/stem/util/system.py
index 14516b2..b9baa28 100644
--- a/stem/util/system.py
+++ b/stem/util/system.py
@@ -27,6 +27,7 @@ best-effort, providing **None** if the lookup fails.
   pid_by_name - gets the pid for a process by the given name
   pid_by_port - gets the pid for a process listening to a given port
   pid_by_open_file - gets the pid for the process with an open file
+  pids_by_user - provides processes owned by a user
   cwd - provides the current working directory for a given process
   user - provides the user a process is running under
   start_time - provides the unix timestamp when the process started
@@ -84,6 +85,7 @@ GET_PID_BY_PORT_NETSTAT = 'netstat -npltu'
 GET_PID_BY_PORT_SOCKSTAT = 'sockstat -4l -P tcp -p %s'
 GET_PID_BY_PORT_LSOF = 'lsof -wnP -iTCP -sTCP:LISTEN'
 GET_PID_BY_FILE_LSOF = 'lsof -tw %s'
+GET_PIDS_BY_USER = 'ps -o pid -u %s'
 GET_CWD_PWDX = 'pwdx %s'
 GET_CWD_LSOF = 'lsof -a -p %s -d cwd -Fn'
 GET_BSD_JAIL_ID_PS = 'ps -p %s -o jid'
@@ -640,6 +642,35 @@ def pid_by_open_file(path):
   return None  # all queries failed
 
 
+def pids_by_user(user):
+  """
+  Provides processes owned by a given user.
+
+  .. versionadded:: 1.5.0
+
+  :param str user: user to look up processes for
+
+  :returns: **list** with the process ids, **None** if it can't be determined
+  """
+
+  # example output:
+  #   atagar at odin:~$ ps -o pid -u avahi
+  #     PID
+  #     914
+  #     915
+
+  if is_available('ps'):
+    results = call(GET_PIDS_BY_USER % user, None)
+
+    if results:
+      try:
+        return list(map(int, results[1:]))
+      except ValueError:
+        pass
+
+  return None
+
+
 def cwd(pid):
   """
   Provides the working directory of the given process.
diff --git a/test/integ/util/system.py b/test/integ/util/system.py
index 1d36b5f..3207f9d 100644
--- a/test/integ/util/system.py
+++ b/test/integ/util/system.py
@@ -366,6 +366,16 @@ class TestSystem(unittest.TestCase):
     os.rmdir(tmpdir)
     self.assertEqual(None, stem.util.system.pid_by_open_file(tmpdir))
 
+  def test_pids_by_user(self):
+    """
+    Checks the stem.util.system.pids_by_user function.
+    """
+
+    # our own pid should be among the processes for our user
+
+    pids = stem.util.system.pids_by_user(getpass.getuser())
+    self.assertTrue(os.getpid() in pids)
+
   def test_cwd(self):
     """
     Checks general usage of the stem.util.system.cwd function.