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 command.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
 161 def get_dir_uid(dir):
 162     """Finds the uid of the person who owns this directory."""
 163     return os.stat(dir).st_uid
 164
 165 def get_dir_owner(dir = "."):
 166     """
 167     Finds the name of the locker this directory is in.
 168
 169     .. note::
 170
 171         This function uses the passwd database and thus
 172         only works on scripts servers when querying directories
 173         that live on AFS.
 174     """
 175     pwentry = pwd.getpwuid(get_dir_uid(dir))
 176     # XXX: Error handling!
 177     return pwentry.pw_name
 178
 179 def get_revision():
 180     """Returns the commit ID of the current Wizard install."""
 181     # If you decide to convert this to use wizard.shell, be warned
 182     # that there is a circular dependency, so this function would
 183     # probably have to live somewhere else, probably wizard.git
 184     wizard_git = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".git")
 185     return subprocess.Popen(["git", "--git-dir=" + wizard_git, "rev-parse", "HEAD"], stdout=subprocess.PIPE).communicate()[0].rstrip()
 186
 187 def get_operator_info():
 188     """
 189     Returns tuple of ``(realname, email)`` about the person running
 190     the script.  If run from a scripts server, get info from Hesiod.
 191     Otherwise, use the passwd database (email generated probably won't
 192     actually accept mail).  Useful when generating commit messages.
 193     """
 194     username = get_operator_name_from_gssapi()
 195     if username:
 196         # scripts approach
 197         hesinfo = subprocess.Popen(["hesinfo", username, "passwd"],stdout=subprocess.PIPE).communicate()[0]
 198         fields = hesinfo.partition(",")[0]
 199         realname = fields.rpartition(":")[2]
 200         return realname, username + "@mit.edu"
 201     else:
 202         # more traditional approach, but the email probably doesn't work
 203         uid = os.getuid()
 204         if not uid:
 205             # since root isn't actually a useful designation, but maybe
 206             # SUDO_USER contains something helpful
 207             sudo_user = os.getenv("SUDO_USER")
 208             if not sudo_user:
 209                 raise NoOperatorInfo
 210             pwdentry = pwd.getpwnam(sudo_user)
 211         else:
 212             pwdentry = pwd.getpwuid(uid)
 213         # XXX: error checking might be nice
 214         # We follow the Ubuntu convention of gecos being a comma split field
 215         # with the person's realname being the first entry.
 216         return pwdentry.pw_gecos.split(",")[0], pwdentry.pw_name + "@" + socket.gethostname()
 217
 218 def get_operator_git():
 219     """
 220     Returns ``Real Name <username@mit.edu>`` suitable for use in
 221     Git ``Something-by:`` string.
 222     """
 223     return "%s <%s>" % get_operator_info()
 224
 225 def get_operator_name_from_gssapi():
 226     """
 227     Returns username of the person operating this script based
 228     off of the :envvar:`SSH_GSSAPI_NAME` environment variable.
 229
 230     .. note::
 231
 232         :envvar:`SSH_GSSAPI_NAME` is not set by a vanilla OpenSSH
 233         distributions.  Scripts servers are patched to support this
 234         environment variable.
 235     """
 236     principal = os.getenv("SSH_GSSAPI_NAME")
 237     if not principal:
 238         return None
 239     instance, _, _ = principal.partition("@")
 240     if instance.endswith("/root"):
 241         username, _, _ = principal.partition("/")
 242     else:
 243         username = instance
 244     return username
 245
 246 def set_operator_env():
 247     """
 248     Sets :envvar:`GIT_COMMITTER_NAME` and :envvar:`GIT_COMMITTER_EMAIL`
 249     environment variables if applicable.  Does nothing if
 250     :func:`get_operator_info` throws :exc:`NoOperatorInfo`.
 251     """
 252     try:
 253         op_realname, op_email = get_operator_info()
 254         os.putenv("GIT_COMMITTER_NAME", op_realname)
 255         os.putenv("GIT_COMMITTER_EMAIL", op_email)
 256     except NoOperatorInfo:
 257         pass
 258
 259 def set_author_env():
 260     """
 261     Sets :envvar:`GIT_AUTHOR_NAME` and :envvar:`GIT_AUTHOR_EMAIL` environment
 262     variables if applicable. Does nothing if :func:`get_dir_owner` fails.
 263     """
 264     try:
 265         # XXX: should check if the directory is in AFS, and if not, use
 266         # a more traditional metric
 267         lockername = get_dir_owner()
 268         os.putenv("GIT_AUTHOR_NAME", "%s locker" % lockername)
 269         os.putenv("GIT_AUTHOR_EMAIL", "%s@scripts.mit.edu" % lockername)
 270     except KeyError: # XXX: This doesn't actually make sense
 271         pass
 272
 273 def set_git_env():
 274     """Sets all appropriate environment variables for Git commits."""
 275     set_operator_env()
 276     set_author_env()
 277
 278 def get_git_footer():
 279     """Returns strings for placing in Git log info about Wizard."""
 280     return "\n".join(["Wizard-revision: %s" % get_revision()
 281         ,"Wizard-args: %s" % " ".join(sys.argv)
 282         ])
 283
 284 def safe_unlink(file):
 285     """Moves a file/dir to a backup location."""
 286     prefix = "%s.bak" % file
 287     name = None
 288     for i in itertools.count():
 289         name = "%s.%d" % (prefix, i)
 290         if not os.path.exists(name):
 291             break
 292     os.rename(file, name)
 293     return name
 294
 295 def fetch(host, path, subpath, post=None):
 296     h = httplib.HTTPConnection(host)
 297     fullpath = path + "/" + subpath
 298     if post:
 299         headers = {"Content-type": "application/x-www-form-urlencoded"}
 300         h.request("POST", fullpath, urllib.urlencode(post), headers)
 301     else:
 302         h.request("GET", fullpath)
 303     r = h.getresponse()
 304     data = r.read()
 305     h.close()
 306     return data
 307
 308 class NoOperatorInfo(wizard.Error):
 309     """No information could be found about the operator from Kerberos."""
 310     pass
 311
 312 class PermissionsError(IOError):
 313     errno = errno.EACCES
 314
 315 class NoSuchDirectoryError(IOError):
 316     errno = errno.ENOENT
 317
 318 class DirectoryLockedError(wizard.Error):
 319     def __init__(self, dir):
 320         self.dir = dir
 321     def __str__(self):
 322         return """
 323
 324 ERROR: Could not acquire lock on directory.  Maybe there is
 325 another migration process running?
 326 """
 327