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 def dictmap(f, d):
  57     """
  58     A map function for dictionaries.  Only changes values.
  59
  60         >>> dictmap(lambda x: x + 2, {'a': 1, 'b': 2})
  61         {'a': 3, 'b': 4}
  62     """
  63     return dict((k,f(v)) for k,v in d.items())
  64
  65 def dictkmap(f, d):
  66     """
  67     A map function for dictionaries that passes key and value.
  68
  69         >>> dictkmap(lambda x, y: x + y, {1: 4, 3: 4})
  70         {1: 5, 3: 6}
  71     """
  72     return dict((k,f(k,v)) for k,v in d.items())
  73
  74 def get_exception_name(output):
  75     """
  76     Reads the traceback from a Python program and grabs the
  77     fully qualified exception name.
  78     """
  79     lines = output.split("\n")
  80     for line in lines[1:]: # skip the "traceback" line
  81         line = line.rstrip()
  82         if line[0] == ' ': continue
  83         if line[-1] == ":":
  84             return line[:-1]
  85         else:
  86             return line
  87
  88 def get_dir_uid(dir):
  89     """Finds the uid of the person who owns this directory."""
  90     return os.stat(dir).st_uid
  91
  92 def get_dir_owner(dir = "."):
  93     """
  94     Finds the name of the locker this directory is in.
  95
  96     .. note::
  97
  98         This function uses the passwd database and thus
  99         only works on scripts servers when querying directories
 100         that live on AFS.
 101     """
 102     pwentry = pwd.getpwuid(get_dir_uid(dir))
 103     # XXX: Error handling!
 104     return pwentry.pw_name
 105
 106 def get_revision():
 107     """Returns the commit ID of the current Wizard install."""
 108     # If you decide to convert this to use wizard.shell, be warned
 109     # that there is a circular dependency, so this function would
 110     # probably have to live somewhere else, probably wizard.git
 111     wizard_git = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".git")
 112     return subprocess.Popen(["git", "--git-dir=" + wizard_git, "rev-parse", "HEAD"], stdout=subprocess.PIPE).communicate()[0].rstrip()
 113
 114 def get_operator_info():
 115     """
 116     Returns tuple of ``(realname, email)`` about the person running
 117     the script.  If run from a scripts server, get info from Hesiod.
 118     Otherwise, use the passwd database (email generated probably won't
 119     actually accept mail).  Useful when generating commit messages.
 120     """
 121     username = get_operator_name_from_gssapi()
 122     if username:
 123         # scripts approach
 124         hesinfo = subprocess.Popen(["hesinfo", username, "passwd"],stdout=subprocess.PIPE).communicate()[0]
 125         fields = hesinfo.partition(",")[0]
 126         realname = fields.rpartition(":")[2]
 127         return realname, username + "@mit.edu"
 128     else:
 129         # more traditional approach, but the email probably doesn't work
 130         uid = os.getuid()
 131         if not uid:
 132             # since root isn't actually a useful designation, but maybe
 133             # SUDO_USER contains something helpful
 134             sudo_user = os.getenv("SUDO_USER")
 135             if not sudo_user:
 136                 raise NoOperatorInfo
 137             pwdentry = pwd.getpwnam(sudo_user)
 138         else:
 139             pwdentry = pwd.getpwuid(uid)
 140         # XXX: error checking might be nice
 141         # We follow the Ubuntu convention of gecos being a comma split field
 142         # with the person's realname being the first entry.
 143         return pwdentry.pw_gecos.split(",")[0], pwdentry.pw_name + "@" + socket.gethostname()
 144
 145 def get_operator_git():
 146     """
 147     Returns ``Real Name <username@mit.edu>`` suitable for use in
 148     Git ``Something-by:`` string.
 149     """
 150     return "%s <%s>" % get_operator_info()
 151
 152 def get_operator_name_from_gssapi():
 153     """
 154     Returns username of the person operating this script based
 155     off of the :envvar:`SSH_GSSAPI_NAME` environment variable.
 156
 157     .. note::
 158
 159         :envvar:`SSH_GSSAPI_NAME` is not set by a vanilla OpenSSH
 160         distributions.  Scripts servers are patched to support this
 161         environment variable.
 162     """
 163     principal = os.getenv("SSH_GSSAPI_NAME")
 164     if not principal:
 165         return None
 166     instance, _, _ = principal.partition("@")
 167     if instance.endswith("/root"):
 168         username, _, _ = principal.partition("/")
 169     else:
 170         username = instance
 171     return username
 172
 173 def set_operator_env():
 174     """
 175     Sets :envvar:`GIT_COMMITTER_NAME` and :envvar:`GIT_COMMITTER_EMAIL`
 176     environment variables if applicable.  Does nothing if
 177     :func:`get_operator_info` throws :exc:`NoOperatorInfo`.
 178     """
 179     try:
 180         op_realname, op_email = get_operator_info()
 181         os.putenv("GIT_COMMITTER_NAME", op_realname)
 182         os.putenv("GIT_COMMITTER_EMAIL", op_email)
 183     except NoOperatorInfo:
 184         pass
 185
 186 def set_author_env():
 187     """
 188     Sets :envvar:`GIT_AUTHOR_NAME` and :envvar:`GIT_AUTHOR_EMAIL` environment
 189     variables if applicable. Does nothing if :func:`get_dir_owner` fails.
 190     """
 191     try:
 192         # XXX: should check if the directory is in AFS, and if not, use
 193         # a more traditional metric
 194         lockername = get_dir_owner()
 195         os.putenv("GIT_AUTHOR_NAME", "%s locker" % lockername)
 196         os.putenv("GIT_AUTHOR_EMAIL", "%s@scripts.mit.edu" % lockername)
 197     except KeyError: # XXX: This doesn't actually make sense
 198         pass
 199
 200 def set_git_env():
 201     """Sets all appropriate environment variables for Git commits."""
 202     set_operator_env()
 203     set_author_env()
 204
 205 def get_git_footer():
 206     """Returns strings for placing in Git log info about Wizard."""
 207     return "\n".join(["Wizard-revision: %s" % get_revision()
 208         ,"Wizard-args: %s" % " ".join(sys.argv)
 209         ])
 210
 211 class NoOperatorInfo(wizard.Error):
 212     """No information could be found about the operator from Kerberos."""
 213     pass
 214