[tor-commits] [stem/master] Adding stem.util.system.get_name_by_pid()

[atagar at torproject.org]

commit 032df08a5591afaf810dabdbca9177fcc528a7ee
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Sep 2 13:16:00 2013 -0700

    Adding stem.util.system.get_name_by_pid()
    
    Function for getting the command running under a given pid. This is the inverse
    of get_pid_by_name().
---
 docs/change_log.rst      |    1 +
 stem/util/system.py      |   42 ++++++++++++++++++++++++++++++++++++++++++
 test/unit/util/system.py |   23 +++++++++++++++++++++++
 3 files changed, 66 insertions(+)

diff --git a/docs/change_log.rst b/docs/change_log.rst
index a368eb0..75a5e81 100644
--- a/docs/change_log.rst
+++ b/docs/change_log.rst
@@ -60,6 +60,7 @@ The following are only available within stem's `git repository
   * :func:`~stem.util.system.set_process_name` inserted spaces between characters (:trac:`8631`)
   * :func:`~stem.util.system.get_pid_by_name` can now pull for all processes with a given name
   * :func:`~stem.util.system.call` ignored the subprocess' exit status
+  * Added :func:`stem.util.system.get_name_by_pid`
   * Added :func:`stem.util.system.get_user`
   * Added :func:`stem.util.system.get_start_time`
   * Added :func:`stem.util.system.get_bsd_jail_path`
diff --git a/stem/util/system.py b/stem/util/system.py
index fcec8d1..01eea06 100644
--- a/stem/util/system.py
+++ b/stem/util/system.py
@@ -16,6 +16,7 @@ best-effort, providing **None** if the lookup fails.
 
   is_available - determines if a command is available on this system
   is_running - determines if a given process is running
+  get_name_by_pid - gets the name for a process by the given pid
   get_pid_by_name - gets the pid for a process by the given name
   get_pid_by_port - gets the pid for a process listening to a given port
   get_pid_by_open_file - gets the pid for the process with an open file
@@ -60,6 +61,7 @@ SHELL_COMMANDS = ['ulimit']
 
 IS_RUNNING_PS_LINUX = "ps -A co command"
 IS_RUNNING_PS_BSD = "ps -ao ucomm="
+GET_NAME_BY_PID_PS = "ps -p %s -o comm"
 GET_PID_BY_NAME_PGREP = "pgrep -x %s"
 GET_PID_BY_NAME_PIDOF = "pidof %s"
 GET_PID_BY_NAME_PS_LINUX = "ps -o pid -C %s"
@@ -225,6 +227,46 @@ def is_running(command):
   return None
 
 
+def get_name_by_pid(pid):
+  """
+  Attempts to determine the name a given process is running under (not
+  including arguments). This uses...
+
+  ::
+
+    1. Information from /proc
+    2. ps -p <pid> -o command
+
+  :param int pid: process id of the process to be queried
+
+  :returns: **str** with the process name, **None** if it can't be determined
+  """
+
+  process_name = None
+
+  if stem.util.proc.is_available():
+    try:
+      process_name = stem.util.proc.get_stats(pid, stem.util.proc.Stat.COMMAND)[0]
+    except IOError:
+      pass
+
+  # attempts to resolve using ps, failing if:
+  # - system's ps variant doesn't handle these flags (none known at the moment)
+  #
+  # example output:
+  #   atagar at morrigan:~$ ps -p 5767 -o comm
+  #   COMMAND
+  #   vim
+
+  if not process_name:
+    results = call(GET_NAME_BY_PID_PS % pid)
+
+    if results and len(results) == 2 and results[0] == 'COMMAND':
+      process_name = results[1].strip()
+
+  return process_name
+
+
 def get_pid_by_name(process_name, multiple = False):
   """
   Attempts to determine the process id for a running process, using...
diff --git a/test/unit/util/system.py b/test/unit/util/system.py
index 628ab71..36204f6 100644
--- a/test/unit/util/system.py
+++ b/test/unit/util/system.py
@@ -134,6 +134,29 @@ class TestSystem(unittest.TestCase):
     self.assertEquals(None, system.is_running("irssi"))
 
   @patch('stem.util.system.call')
+  @patch('stem.util.proc.is_available', Mock(return_value = False))
+  @patch('stem.util.system.is_available', Mock(return_value = True))
+  def test_get_name_by_pid_ps(self, call_mock):
+    """
+    Tests the get_name_by_pid function with ps responses.
+    """
+
+    responses = {
+      "success": ["COMMAND", "vim"],
+      "malformed_command_1": ["COMMAND"],
+      "malformed_command_2": ["foobar"],
+      "malformed_command_3": ["NOT_COMMAND", "vim"],
+      "no_results": [],
+      "command_fails": None,
+    }
+
+    call_mock.side_effect = mock_call(system.GET_NAME_BY_PID_PS, responses)
+
+    for test_input in responses:
+      expected_response = "vim" if test_input == "success" else None
+      self.assertEquals(expected_response, system.get_name_by_pid(test_input))
+
+  @patch('stem.util.system.call')
   @patch('stem.util.system.is_available', Mock(return_value = True))
   def test_get_pid_by_name_pgrep(self, call_mock):
     """