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 import errno
  16 import itertools
  17 import signal
  18 import httplib
  19 import urllib
  20
  21 import wizard
  22
  23 class ChangeDirectory(object):
  24     """
  25     Context for temporarily changing the working directory.
  26
  27         >>> with ChangeDirectory("/tmp"):
  28         ...    print os.getcwd()
  29         /tmp
  30     """
  31     def __init__(self, dir):
  32         self.dir = dir
  33         self.olddir = None
  34     def __enter__(self):
  35         self.olddir = os.getcwd()
  36         chdir(self.dir)
  37     def __exit__(self, *args):
  38         chdir(self.olddir)
  39
  40 class Counter(object):
  41     """
  42     Object for counting different values when you don't know what
  43     they are a priori.  Supports index access and iteration.
  44
  45         >>> counter = Counter()
  46         >>> counter.count("foo")
  47         >>> print counter["foo"]
  48         1
  49     """
  50     def __init__(self):
  51         self.dict = {}
  52     def count(self, value):
  53         """Increments count for ``value``."""
  54         self.dict.setdefault(value, 0)
  55         self.dict[value] += 1
  56     def __getitem__(self, key):
  57         return self.dict[key]
  58     def __iter__(self):
  59         return self.dict.__iter__()
  60
  61 class PipeToLess(object):
  62     """
  63     Context for printing output to a pager.  Use this if output
  64     is expected to be long.
  65     """
  66     def __enter__(self):
  67         self.proc = subprocess.Popen("less", stdin=subprocess.PIPE)
  68         self.old_stdout = sys.stdout
  69         sys.stdout = self.proc.stdin
  70     def __exit__(self, *args):
  71         if self.proc:
  72             self.proc.stdin.close()
  73             self.proc.wait()
  74             sys.stdout = self.old_stdout
  75
  76 class IgnoreKeyboardInterrupts(object):
  77     """
  78     Context for temporarily ignoring keyboard interrupts.  Use this
  79     if aborting would cause more harm than finishing the job.
  80     """
  81     def __enter__(self):
  82         signal.signal(signal.SIGINT,signal.SIG_IGN)
  83     def __exit__(self, *args):
  84         signal.signal(signal.SIGINT, signal.default_int_handler)
  85
  86 class LockDirectory(object):
  87     """
  88     Context for locking a directory.
  89     """
  90     def __init__(self, lockfile):
  91         self.lockfile = lockfile
  92     def __enter__(self):
  93         try:
  94             os.open(self.lockfile, os.O_CREAT | os.O_EXCL)
  95         except OSError as e:
  96             if e.errno == errno.EEXIST:
  97                 raise DirectoryLockedError(os.getcwd())
  98             elif e.errno == errno.EACCES:
  99                 raise PermissionsError(os.getcwd())
 100             raise
 101     def __exit__(self, *args):
 102         try:
 103             os.unlink(self.lockfile)
 104         except OSError:
 105             pass
 106
 107 def chdir(dir):
 108     """
 109     Changes a directory, but has special exceptions for certain
 110     classes of errors.
 111     """
 112     try:
 113         os.chdir(dir)
 114     except OSError as e:
 115         if e.errno == errno.EACCES:
 116             raise PermissionsError()
 117         elif e.errno == errno.ENOENT:
 118             raise NoSuchDirectoryError()
 119         else: raise e
 120
 121 def dictmap(f, d):
 122     """
 123     A map function for dictionaries.  Only changes values.
 124
 125         >>> dictmap(lambda x: x + 2, {'a': 1, 'b': 2})
 126         {'a': 3, 'b': 4}
 127     """
 128     return dict((k,f(v)) for k,v in d.items())
 129
 130 def dictkmap(f, d):
 131     """
 132     A map function for dictionaries that passes key and value.
 133
 134         >>> dictkmap(lambda x, y: x + y, {1: 4, 3: 4})
 135         {1: 5, 3: 7}
 136     """
 137     return dict((k,f(k,v)) for k,v in d.items())
 138
 139 def get_exception_name(output):
 140     """
 141     Reads the traceback from a Python program and grabs the
 142     fully qualified exception name.
 143     """
 144     lines = output.split("\n")
 145     cue = False
 146     result = "(unknown)"
 147     for line in lines[1:]:
 148         line = line.rstrip()
 149         if not line: continue
 150         if line[0] == ' ':
 151             cue = True
 152             continue
 153         if cue:
 154             cue = False
 155             if line[-1] == ":":
 156                 result = line[:-1]
 157             else:
 158                 result = line
 159             return result
 160     return result
 161
 162 def get_dir_uid(dir):
 163     """Finds the uid of the person who owns this directory."""
 164     return os.stat(dir).st_uid
 165
 166 def get_dir_owner(dir = "."):
 167     """
 168     Finds the name of the locker this directory is in.
 169
 170     .. note::
 171
 172         This function uses the passwd database and thus
 173         only works on scripts servers when querying directories
 174         that live on AFS.
 175     """
 176     pwentry = pwd.getpwuid(get_dir_uid(dir))
 177     # XXX: Error handling!
 178     return pwentry.pw_name
 179
 180 def get_revision():
 181     """Returns the commit ID of the current Wizard install."""
 182     # If you decide to convert this to use wizard.shell, be warned
 183     # that there is a circular dependency, so this function would
 184     # probably have to live somewhere else, probably wizard.git
 185     wizard_git = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".git")
 186     return subprocess.Popen(["git", "--git-dir=" + wizard_git, "rev-parse", "HEAD"], stdout=subprocess.PIPE).communicate()[0].rstrip()
 187
 188 def get_operator_info():
 189     """
 190     Returns tuple of ``(realname, email)`` about the person running
 191     the script.  If run from a scripts server, get info from Hesiod.
 192     Otherwise, use the passwd database (email generated probably won't
 193     actually accept mail).  Useful when generating commit messages.
 194     """
 195     username = get_operator_name_from_gssapi()
 196     if username:
 197         # scripts approach
 198         hesinfo = subprocess.Popen(["hesinfo", username, "passwd"],stdout=subprocess.PIPE).communicate()[0]
 199         fields = hesinfo.partition(",")[0]
 200         realname = fields.rpartition(":")[2]
 201         return realname, username + "@mit.edu"
 202     else:
 203         # more traditional approach, but the email probably doesn't work
 204         uid = os.getuid()
 205         if not uid:
 206             # since root isn't actually a useful designation, but maybe
 207             # SUDO_USER contains something helpful
 208             sudo_user = os.getenv("SUDO_USER")
 209             if not sudo_user:
 210                 raise NoOperatorInfo
 211             pwdentry = pwd.getpwnam(sudo_user)
 212         else:
 213             pwdentry = pwd.getpwuid(uid)
 214         # XXX: error checking might be nice
 215         # We follow the Ubuntu convention of gecos being a comma split field
 216         # with the person's realname being the first entry.
 217         return pwdentry.pw_gecos.split(",")[0], pwdentry.pw_name + "@" + socket.gethostname()
 218
 219 def get_operator_git():
 220     """
 221     Returns ``Real Name <username@mit.edu>`` suitable for use in
 222     Git ``Something-by:`` string.
 223     """
 224     return "%s <%s>" % get_operator_info()
 225
 226 def get_operator_name_from_gssapi():
 227     """
 228     Returns username of the person operating this script based
 229     off of the :envvar:`SSH_GSSAPI_NAME` environment variable.
 230
 231     .. note::
 232
 233         :envvar:`SSH_GSSAPI_NAME` is not set by a vanilla OpenSSH
 234         distributions.  Scripts servers are patched to support this
 235         environment variable.
 236     """
 237     principal = os.getenv("SSH_GSSAPI_NAME")
 238     if not principal:
 239         return None
 240     instance, _, _ = principal.partition("@")
 241     if instance.endswith("/root"):
 242         username, _, _ = principal.partition("/")
 243     else:
 244         username = instance
 245     return username
 246
 247 def set_operator_env():
 248     """
 249     Sets :envvar:`GIT_COMMITTER_NAME` and :envvar:`GIT_COMMITTER_EMAIL`
 250     environment variables if applicable.  Does nothing if
 251     :func:`get_operator_info` throws :exc:`NoOperatorInfo`.
 252     """
 253     try:
 254         op_realname, op_email = get_operator_info()
 255         os.putenv("GIT_COMMITTER_NAME", op_realname)
 256         os.putenv("GIT_COMMITTER_EMAIL", op_email)
 257     except NoOperatorInfo:
 258         pass
 259
 260 def set_author_env():
 261     """
 262     Sets :envvar:`GIT_AUTHOR_NAME` and :envvar:`GIT_AUTHOR_EMAIL` environment
 263     variables if applicable. Does nothing if :func:`get_dir_owner` fails.
 264     """
 265     try:
 266         # XXX: should check if the directory is in AFS, and if not, use
 267         # a more traditional metric
 268         lockername = get_dir_owner()
 269         os.putenv("GIT_AUTHOR_NAME", "%s locker" % lockername)
 270         os.putenv("GIT_AUTHOR_EMAIL", "%s@scripts.mit.edu" % lockername)
 271     except KeyError: # XXX: This doesn't actually make sense
 272         pass
 273
 274 def set_git_env():
 275     """Sets all appropriate environment variables for Git commits."""
 276     set_operator_env()
 277     set_author_env()
 278
 279 def get_git_footer():
 280     """Returns strings for placing in Git log info about Wizard."""
 281     return "\n".join(["Wizard-revision: %s" % get_revision()
 282         ,"Wizard-args: %s" % " ".join(sys.argv)
 283         ])
 284
 285 def safe_unlink(file):
 286     """Moves a file/dir to a backup location."""
 287     if not os.path.exists(file):
 288         return None
 289     prefix = "%s.bak" % file
 290     name = None
 291     for i in itertools.count():
 292         name = "%s.%d" % (prefix, i)
 293         if not os.path.exists(name):
 294             break
 295     os.rename(file, name)
 296     return name
 297
 298 def fetch(host, path, subpath, post=None):
 299     h = httplib.HTTPConnection(host)
 300     fullpath = path + "/" + subpath
 301     if post:
 302         headers = {"Content-type": "application/x-www-form-urlencoded"}
 303         h.request("POST", fullpath, urllib.urlencode(post), headers)
 304     else:
 305         h.request("GET", fullpath)
 306     r = h.getresponse()
 307     data = r.read()
 308     h.close()
 309     return data
 310
 311 class NoOperatorInfo(wizard.Error):
 312     """No information could be found about the operator from Kerberos."""
 313     pass
 314
 315 class PermissionsError(IOError):
 316     errno = errno.EACCES
 317
 318 class NoSuchDirectoryError(IOError):
 319     errno = errno.ENOENT
 320
 321 class DirectoryLockedError(wizard.Error):
 322     def __init__(self, dir):
 323         self.dir = dir
 324     def __str__(self):
 325         return """
 326
 327 ERROR: Could not acquire lock on directory.  Maybe there is
 328 another migration process running?
 329 """
 330