[tor-commits] [arm/master] Renaming our process to "arm <input args>"

[atagar at torproject.org]

commit 8ec467eb5caeb4972e12a988a805d93f511bd192
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon May 2 19:59:39 2011 -0700

    Renaming our process to "arm <input args>"
    
    Our process was called "python ./src/starter.py <input args>" which made it a
    pita to identify, use killall on, etc. Using some ctypes hacks to overwrite it
    with an intuitive name.
---
 src/starter.py       |    8 +++
 src/util/procname.py |  122 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 130 insertions(+), 0 deletions(-)

diff --git a/src/starter.py b/src/starter.py
index c0a0270..365676c 100644
--- a/src/starter.py
+++ b/src/starter.py
@@ -387,5 +387,13 @@ if __name__ == '__main__':
     util.log.log(CONFIG["log.savingDebugLog"], "Saving a debug log to '%s' (please check it for sensitive information before sharing)" % LOG_DUMP_PATH)
     _dumpConfig()
   
+  # Attempts to rename our process from "python setup.py <input args>" to
+  # "arm <input args>"
+  
+  try:
+    from util import procname
+    procname.renameProcess("arm %s" % " ".join(sys.argv[1:]))
+  except: pass
+  
   cli.controller.startTorMonitor(time.time() - initTime, expandedEvents, param["startup.blindModeEnabled"])
 
diff --git a/src/util/procname.py b/src/util/procname.py
new file mode 100644
index 0000000..7641c5a
--- /dev/null
+++ b/src/util/procname.py
@@ -0,0 +1,122 @@
+# Module to allow for arbitrary renaming of our python process. This is mostly
+# based on:
+# http://www.rhinocerus.net/forum/lang-python/569677-setting-program-name-like-0-perl.html#post2272369
+# and an adaptation by Jake: https://github.com/ioerror/chameleon
+# 
+# A cleaner implementation is available at:
+# https://github.com/cream/libs/blob/b38970e2a6f6d2620724c828808235be0445b799/cream/util/procname.py
+# but I'm not quite clear on their implementation, and it only does targeted
+# argument replacement (ie, replace argv[0], argv[1], etc but with a string
+# the same size).
+
+import os
+import sys
+import ctypes
+import ctypes.util
+
+from util import sysTools
+
+# flag for setting the process name, found in '/usr/include/linux/prctl.h'
+PR_SET_NAME = 15
+
+# Maximum number of characters we'll set the process name to. Evidently this
+# cap was simply chosen since it didn't cause a segfault for its author.
+MAX_CHAR = 1608
+
+argc_t = ctypes.POINTER(ctypes.c_char_p)
+
+Py_GetArgcArgv = ctypes.pythonapi.Py_GetArgcArgv
+Py_GetArgcArgv.restype = None
+Py_GetArgcArgv.argtypes = [ctypes.POINTER(ctypes.c_int),
+                           ctypes.POINTER(argc_t)]
+
+# tracks the last name we've changed the process to
+currentProcessName = None
+
+def renameProcess(processName):
+  """
+  Renames our current process from "python <args>" to a custom name.
+  
+  Arguments:
+    processName - new name for our process
+  """
+  
+  _setArgv(processName)
+  if sys.platform == "linux2":
+    _setPrctlName(processName)
+  elif sys.platform == "freebsd7":
+    _setProcTitle(processName)
+
+def _setArgv(processName):
+  """
+  Overwrites our argv in a similar fashion to how it's done in C with:
+  strcpy(argv[0], "new_name");
+  """
+  
+  global currentProcessName
+  
+  argv = ctypes.c_int(0)
+  argc = argc_t()
+  Py_GetArgcArgv(argv, ctypes.pointer(argc))
+  
+  # The original author did the memset for 256, while Jake did it for the
+  # processName length (capped at 1608). I'm not sure of the reasons for
+  # either of these limits, but setting it to anything higher than than the
+  # length of the null terminated process name should be pointless, so opting
+  # for Jake's implementation on this.
+  
+  if currentProcessName == None:
+    # Using argv via...
+    # currentProcessName = " ".join(["python"] + sys.argv)
+    #
+    # doesn't do the trick since this will miss interpretor arguments like...
+    # python -W ignore::DeprecationWarning myScript.py
+    # 
+    # hence simply getting an outside opinion of our command.
+    
+    psResults = sysTools.call("ps -p %i -o args" % os.getpid())
+    
+    if len(psResults) == 2:
+      # output looks like:
+      # COMMAND
+      # python ./src/starter.py
+      
+      currentProcessName = psResults[1]
+    
+    if not currentProcessName:
+      raise IOError("unable to determine our process name")
+  
+  # we need to fill extra space with null characters, otherwise the process
+  # will end with a '?' when you run ps against it
+  if len(currentProcessName) > len(processName):
+    processName += "\0" * (len(currentProcessName) - len(processName))
+  
+  size = min(len(processName), MAX_CHAR)
+  ctypes.memset(argc.contents, 0, size + 1) # null terminate the string's end
+  ctypes.memmove(argc.contents, processName, size)
+  currentProcessName = processName[:MAX_CHAR]
+
+def _setPrctlName(processName):
+  """
+  Sets the prctl name, which is used by top and killall. This appears to be
+  Linux specific and has the max of 15 characters. Source:
+  http://stackoverflow.com/questions/564695/is-there-a-way-to-change-effective-process-name-in-python/923034#923034
+  """
+  
+  libc = ctypes.CDLL(ctypes.util.find_library("c"))
+  nameBuffer = ctypes.create_string_buffer(len(processName)+1)
+  nameBuffer.value = processName
+  libc.prctl(PR_SET_NAME, ctypes.byref(nameBuffer), 0, 0, 0)
+
+def _setProcTitle(processName):
+  """
+  BSD specific calls (should be compataible with both FreeBSD and OpenBSD:
+  http://fxr.watson.org/fxr/source/gen/setproctitle.c?v=FREEBSD-LIBC
+  http://www.rootr.net/man/man/setproctitle/3
+  """
+  
+  libc = ctypes.CDLL(ctypes.util.find_library("c"))
+  nameBuffer = ctypes.create_string_buffer(len(processName)+1)
+  nameBuffer.value = processName
+  libc.setproctitle(ctypes.byref(nameBuffer))
+