wizard/util.py

   1 """
   2 Miscellaneous utility functions and classes.
   3
   4 .. testsetup:: *
   5
   6     from wizard.util import *
   7 """
   8
   9 import os.path
  10 import os
  11 import subprocess
  12 import pwd
  13 import sys
  14 import socket
  15
  16 import wizard
  17
  18 class ChangeDirectory(object):
  19     """
  20     Context for temporarily changing the working directory.
  21
  22         >>> with ChangeDirectory("/tmp"):
  23         ...    print os.getcwd()
  24         /tmp
  25     """
  26     def __init__(self, dir):
  27         self.dir = dir
  28         self.olddir = None
  29     def __enter__(self):
  30         self.olddir = os.getcwd()
  31         os.chdir(self.dir)
  32     def __exit__(self, *args):
  33         os.chdir(self.olddir)
  34
  35 class Counter(object):
  36     """
  37     Object for counting different values when you don't know what
  38     they are a priori.  Supports index access and iteration.
  39
  40         >>> counter = Counter()
  41         >>> counter.count("foo")
  42         >>> print counter["foo"]
  43         1
  44     """
  45     def __init__(self):
  46         self.dict = {}
  47     def count(self, value):
  48         """Increments count for ``value``."""
  49         self.dict.setdefault(value, 0)
  50         self.dict[value] += 1
  51     def __getitem__(self, key):
  52         return self.dict[key]
  53     def __iter__(self):
  54         return self.dict.__iter__()
  55
  56 class PipeToLess(object):
  57     """
  58     Context for printing output to a pager.  Use this if output
  59     is expected to be long.
  60     """
  61     def __enter__(self):
  62         self.proc = subprocess.Popen("less", stdin=subprocess.PIPE)
  63         self.old_stdout = sys.stdout
  64         sys.stdout = self.proc.stdin
  65     def __exit__(self, *args):
  66         if self.proc:
  67             self.proc.stdin.close()
  68             self.proc.wait()
  69             sys.stdout = self.old_stdout
  70
  71 def dictmap(f, d):
  72     """
  73     A map function for dictionaries.  Only changes values.
  74
  75         >>> dictmap(lambda x: x + 2, {'a': 1, 'b': 2})
  76         {'a': 3, 'b': 4}
  77     """
  78     return dict((k,f(v)) for k,v in d.items())
  79
  80 def dictkmap(f, d):
  81     """
  82     A map function for dictionaries that passes key and value.
  83
  84         >>> dictkmap(lambda x, y: x + y, {1: 4, 3: 4})
  85         {1: 5, 3: 6}
  86     """
  87     return dict((k,f(k,v)) for k,v in d.items())
  88
  89 def get_exception_name(output):
  90     """
  91     Reads the traceback from a Python program and grabs the
  92     fully qualified exception name.
  93     """
  94     lines = output.split("\n")
  95     cue = False
  96     for line in lines[1:]:
  97         line = line.rstrip()
  98         if not line: continue
  99         if line[0] == ' ':
 100             cue = True
 101             continue
 102         if cue:
 103             cue = False
 104             if line[-1] == ":":
 105                 result = line[:-1]
 106             else:
 107                 result = line
 108     return result
 109
 110 def get_dir_uid(dir):
 111     """Finds the uid of the person who owns this directory."""
 112     return os.stat(dir).st_uid
 113
 114 def get_dir_owner(dir = "."):
 115     """
 116     Finds the name of the locker this directory is in.
 117
 118     .. note::
 119
 120         This function uses the passwd database and thus
 121         only works on scripts servers when querying directories
 122         that live on AFS.
 123     """
 124     pwentry = pwd.getpwuid(get_dir_uid(dir))
 125     # XXX: Error handling!
 126     return pwentry.pw_name
 127
 128 def get_revision():
 129     """Returns the commit ID of the current Wizard install."""
 130     # If you decide to convert this to use wizard.shell, be warned
 131     # that there is a circular dependency, so this function would
 132     # probably have to live somewhere else, probably wizard.git
 133     wizard_git = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".git")
 134     return subprocess.Popen(["git", "--git-dir=" + wizard_git, "rev-parse", "HEAD"], stdout=subprocess.PIPE).communicate()[0].rstrip()
 135
 136 def get_operator_info():
 137     """
 138     Returns tuple of ``(realname, email)`` about the person running
 139     the script.  If run from a scripts server, get info from Hesiod.
 140     Otherwise, use the passwd database (email generated probably won't
 141     actually accept mail).  Useful when generating commit messages.
 142     """
 143     username = get_operator_name_from_gssapi()
 144     if username:
 145         # scripts approach
 146         hesinfo = subprocess.Popen(["hesinfo", username, "passwd"],stdout=subprocess.PIPE).communicate()[0]
 147         fields = hesinfo.partition(",")[0]
 148         realname = fields.rpartition(":")[2]
 149         return realname, username + "@mit.edu"
 150     else:
 151         # more traditional approach, but the email probably doesn't work
 152         uid = os.getuid()
 153         if not uid:
 154             # since root isn't actually a useful designation, but maybe
 155             # SUDO_USER contains something helpful
 156             sudo_user = os.getenv("SUDO_USER")
 157             if not sudo_user:
 158                 raise NoOperatorInfo
 159             pwdentry = pwd.getpwnam(sudo_user)
 160         else:
 161             pwdentry = pwd.getpwuid(uid)
 162         # XXX: error checking might be nice
 163         # We follow the Ubuntu convention of gecos being a comma split field
 164         # with the person's realname being the first entry.
 165         return pwdentry.pw_gecos.split(",")[0], pwdentry.pw_name + "@" + socket.gethostname()
 166
 167 def get_operator_git():
 168     """
 169     Returns ``Real Name <username@mit.edu>`` suitable for use in
 170     Git ``Something-by:`` string.
 171     """
 172     return "%s <%s>" % get_operator_info()
 173
 174 def get_operator_name_from_gssapi():
 175     """
 176     Returns username of the person operating this script based
 177     off of the :envvar:`SSH_GSSAPI_NAME` environment variable.
 178
 179     .. note::
 180
 181         :envvar:`SSH_GSSAPI_NAME` is not set by a vanilla OpenSSH
 182         distributions.  Scripts servers are patched to support this
 183         environment variable.
 184     """
 185     principal = os.getenv("SSH_GSSAPI_NAME")
 186     if not principal:
 187         return None
 188     instance, _, _ = principal.partition("@")
 189     if instance.endswith("/root"):
 190         username, _, _ = principal.partition("/")
 191     else:
 192         username = instance
 193     return username
 194
 195 def set_operator_env():
 196     """
 197     Sets :envvar:`GIT_COMMITTER_NAME` and :envvar:`GIT_COMMITTER_EMAIL`
 198     environment variables if applicable.  Does nothing if
 199     :func:`get_operator_info` throws :exc:`NoOperatorInfo`.
 200     """
 201     try:
 202         op_realname, op_email = get_operator_info()
 203         os.putenv("GIT_COMMITTER_NAME", op_realname)
 204         os.putenv("GIT_COMMITTER_EMAIL", op_email)
 205     except NoOperatorInfo:
 206         pass
 207
 208 def set_author_env():
 209     """
 210     Sets :envvar:`GIT_AUTHOR_NAME` and :envvar:`GIT_AUTHOR_EMAIL` environment
 211     variables if applicable. Does nothing if :func:`get_dir_owner` fails.
 212     """
 213     try:
 214         # XXX: should check if the directory is in AFS, and if not, use
 215         # a more traditional metric
 216         lockername = get_dir_owner()
 217         os.putenv("GIT_AUTHOR_NAME", "%s locker" % lockername)
 218         os.putenv("GIT_AUTHOR_EMAIL", "%s@scripts.mit.edu" % lockername)
 219     except KeyError: # XXX: This doesn't actually make sense
 220         pass
 221
 222 def set_git_env():
 223     """Sets all appropriate environment variables for Git commits."""
 224     set_operator_env()
 225     set_author_env()
 226
 227 def get_git_footer():
 228     """Returns strings for placing in Git log info about Wizard."""
 229     return "\n".join(["Wizard-revision: %s" % get_revision()
 230         ,"Wizard-args: %s" % " ".join(sys.argv)
 231         ])
 232
 233 class NoOperatorInfo(wizard.Error):
 234     """No information could be found about the operator from Kerberos."""
 235     pass
 236