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 import time
  21 import logging
  22 import random
  23 import string
  24
  25 import wizard
  26 from wizard import user
  27
  28 def boolish(val):
  29     """
  30     Parse the contents of an environment variable as a boolean.
  31     This recognizes more values as ``False`` than :func:`bool` would.
  32
  33         >>> boolish("0")
  34         False
  35         >>> boolish("no")
  36         False
  37         >>> boolish("1")
  38         True
  39     """
  40     try:
  41         return bool(int(val))
  42     except (ValueError, TypeError):
  43         if val == "No" or val == "no" or val == "false" or val == "False":
  44             return False
  45         return bool(val)
  46
  47 class ChangeDirectory(object):
  48     """
  49     Context for temporarily changing the working directory.
  50
  51         >>> with ChangeDirectory("/tmp"):
  52         ...    print os.getcwd()
  53         /tmp
  54     """
  55     def __init__(self, dir):
  56         self.dir = dir
  57         self.olddir = None
  58     def __enter__(self):
  59         self.olddir = os.getcwd()
  60         chdir(self.dir)
  61     def __exit__(self, *args):
  62         chdir(self.olddir)
  63
  64 class Counter(object):
  65     """
  66     Object for counting different values when you don't know what
  67     they are a priori.  Supports index access and iteration.
  68
  69         >>> counter = Counter()
  70         >>> counter.count("foo")
  71         >>> print counter["foo"]
  72         1
  73     """
  74     def __init__(self):
  75         self.dict = {}
  76     def count(self, value):
  77         """Increments count for ``value``."""
  78         self.dict.setdefault(value, 0)
  79         self.dict[value] += 1
  80     def __getitem__(self, key):
  81         return self.dict[key]
  82     def __iter__(self):
  83         return self.dict.__iter__()
  84     def max(self):
  85         """Returns the max counter value seen."""
  86         return max(self.dict.values())
  87     def sum(self):
  88         """Returns the sum of all counter values."""
  89         return sum(self.dict.values())
  90     def keys(self):
  91         """Returns the keys of counters."""
  92         return self.dict.keys()
  93
  94 class PipeToLess(object):
  95     """
  96     Context for printing output to a pager.  Use this if output
  97     is expected to be long.
  98     """
  99     def __enter__(self):
 100         self.proc = subprocess.Popen("less", stdin=subprocess.PIPE)
 101         self.old_stdout = sys.stdout
 102         sys.stdout = self.proc.stdin
 103     def __exit__(self, *args):
 104         if self.proc:
 105             self.proc.stdin.close()
 106             self.proc.wait()
 107             sys.stdout = self.old_stdout
 108
 109 class IgnoreKeyboardInterrupts(object):
 110     """
 111     Context for temporarily ignoring keyboard interrupts.  Use this
 112     if aborting would cause more harm than finishing the job.
 113     """
 114     def __enter__(self):
 115         signal.signal(signal.SIGINT,signal.SIG_IGN)
 116     def __exit__(self, *args):
 117         signal.signal(signal.SIGINT, signal.default_int_handler)
 118
 119 class LockDirectory(object):
 120     """
 121     Context for locking a directory.
 122     """
 123     def __init__(self, lockfile, expiry = 3600):
 124         self.lockfile = lockfile
 125         self.expiry = expiry # by default an hour
 126     def __enter__(self):
 127         # It's A WAVY
 128         for i in range(0, 3):
 129             try:
 130                 os.open(self.lockfile, os.O_CREAT | os.O_EXCL)
 131                 open(self.lockfile, "w").write("%d" % os.getpid())
 132             except OSError as e:
 133                 if e.errno == errno.EEXIST:
 134                     # There is a possibility of infinite recursion, but we
 135                     # expect it to be unlikely, and not harmful if it does happen
 136                     with LockDirectory(self.lockfile + "_"):
 137                         # See if we can break the lock
 138                         try:
 139                             pid = open(self.lockfile, "r").read().strip()
 140                             if not os.path.exists("/proc/%s" % pid):
 141                                 # break the lock, try again
 142                                 logging.warning("Breaking orphaned lock at %s", self.lockfile)
 143                                 os.unlink(self.lockfile)
 144                                 continue
 145                             try:
 146                                 # check if the file is expiry old, if so, break the lock, try again
 147                                 if time.time() - os.stat(self.lockfile).st_mtime > self.expiry:
 148                                     logging.warning("Breaking stale lock at %s", self.lockfile)
 149                                     os.unlink(self.lockfile)
 150                                     continue
 151                             except OSError as e:
 152                                 if e.errno == errno.ENOENT:
 153                                     continue
 154                                 raise
 155                         except IOError:
 156                             # oh hey, it went away; try again
 157                             continue
 158                     raise DirectoryLockedError(os.getcwd())
 159                 elif e.errno == errno.EACCES:
 160                     raise PermissionsError(os.getcwd())
 161                 raise
 162             return
 163         raise DirectoryLockedError(os.getcwd())
 164     def __exit__(self, *args):
 165         try:
 166             os.unlink(self.lockfile)
 167         except OSError:
 168             pass
 169
 170 def chdir(dir):
 171     """
 172     Changes a directory, but has special exceptions for certain
 173     classes of errors.
 174     """
 175     try:
 176         os.chdir(dir)
 177     except OSError as e:
 178         if e.errno == errno.EACCES:
 179             raise PermissionsError()
 180         elif e.errno == errno.ENOENT:
 181             raise NoSuchDirectoryError()
 182         else: raise e
 183
 184 def dictmap(f, d):
 185     """
 186     A map function for dictionaries.  Only changes values.
 187
 188         >>> dictmap(lambda x: x + 2, {'a': 1, 'b': 2})
 189         {'a': 3, 'b': 4}
 190     """
 191     return dict((k,f(v)) for k,v in d.items())
 192
 193 def dictkmap(f, d):
 194     """
 195     A map function for dictionaries that passes key and value.
 196
 197         >>> dictkmap(lambda x, y: x + y, {1: 4, 3: 4})
 198         {1: 5, 3: 7}
 199     """
 200     return dict((k,f(k,v)) for k,v in d.items())
 201
 202 def get_exception_name(output):
 203     """
 204     Reads the traceback from a Python program and grabs the
 205     fully qualified exception name.
 206     """
 207     lines = output.split("\n")
 208     cue = False
 209     result = "(unknown)"
 210     for line in lines[1:]:
 211         line = line.rstrip()
 212         if not line: continue
 213         if line[0] == ' ':
 214             cue = True
 215             continue
 216         if cue:
 217             cue = False
 218             return line.partition(':')[0]
 219     return result
 220
 221 def get_dir_uid(dir):
 222     """Finds the uid of the person who owns this directory."""
 223     return os.stat(dir).st_uid
 224
 225 def get_revision():
 226     """Returns the commit ID of the current Wizard install."""
 227     # If you decide to convert this to use wizard.shell, be warned
 228     # that there is a circular dependency, so this function would
 229     # probably have to live somewhere else, probably wizard.git
 230     wizard_git = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".git")
 231     return subprocess.Popen(["git", "--git-dir=" + wizard_git, "rev-parse", "HEAD"], stdout=subprocess.PIPE).communicate()[0].rstrip()
 232
 233 def get_operator_git():
 234     """
 235     Returns ``Real Name <username@mit.edu>`` suitable for use in
 236     Git ``Something-by:`` string.  Throws :exc:`NoOperatorInfo` if
 237     no operator information is available.
 238     """
 239     op = user.operator()
 240     if op is None:
 241         raise NoOperatorInfo
 242     info = user.pwnam(op)
 243     return "%s <%s>" % (info.realname, info.email)
 244
 245 def set_operator_env():
 246     """
 247     Sets :envvar:`GIT_COMMITTER_NAME` and :envvar:`GIT_COMMITTER_EMAIL`
 248     environment variables if applicable.  Does nothing if no information
 249     is available
 250     """
 251     op = user.operator()
 252     if op is None:
 253         return
 254     info = user.pwnam(op)
 255     if not info.realname:
 256         return
 257     os.putenv("GIT_COMMITTER_NAME", info.realname)
 258     os.putenv("GIT_COMMITTER_EMAIL", info.email)
 259
 260 def set_author_env():
 261     """
 262     Sets :envvar:`GIT_AUTHOR_NAME` and :envvar:`GIT_AUTHOR_EMAIL`
 263     environment variables if applicable. Does nothing if
 264     :func:`wizard.user.passwd` fails.
 265     """
 266     info = user.passwd()
 267     if info is None:
 268         return
 269     if not info.realname:
 270         return
 271     os.putenv("GIT_AUTHOR_NAME", "%s" % info.realname)
 272     os.putenv("GIT_AUTHOR_EMAIL", "%s" % info.email)
 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.lexists(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.lexists(name):
 294             break
 295     os.rename(file, name)
 296     return name
 297
 298 def soft_unlink(file):
 299     """Unlink a file, but don't complain if it doesn't exist."""
 300     try:
 301         os.unlink(file)
 302     except OSError:
 303         pass
 304
 305 def makedirs(path):
 306     """
 307     Create a directory path (a la ``mkdir -p`` or ``os.makedirs``),
 308     but don't complain if it already exists.
 309     """
 310     try:
 311         os.makedirs(path)
 312     except OSError as exc:
 313         if exc.errno == errno.EEXIST:
 314             pass
 315         else:
 316             raise
 317
 318 def fetch(host, path, subpath, post=None):
 319     try:
 320         # XXX: Should use urllib instead
 321         h = httplib.HTTPConnection(host)
 322         fullpath = path.rstrip("/") + "/" + subpath.lstrip("/") # to be lenient about input we accept
 323         if post:
 324             headers = {"Content-type": "application/x-www-form-urlencoded"}
 325             logging.info("POST request to http://%s%s", host, fullpath)
 326             logging.debug("POST contents:\n" + urllib.urlencode(post))
 327             h.request("POST", fullpath, urllib.urlencode(post), headers)
 328         else:
 329             logging.info("GET request to http://%s%s", host, fullpath)
 330             h.request("GET", fullpath)
 331         r = h.getresponse()
 332         data = r.read()
 333         h.close()
 334         return data
 335     except socket.gaierror as e:
 336         if e.errno == socket.EAI_NONAME:
 337             raise DNSError(host)
 338         else:
 339             raise
 340
 341 def mixed_newlines(filename):
 342     """Returns ``True`` if ``filename`` has mixed newlines."""
 343     f = open(filename, "U") # requires universal newline support
 344     f.read()
 345     ret = isinstance(f.newlines, tuple)
 346     f.close() # just to be safe
 347     return ret
 348
 349 def disk_usage(dir=None, excluded_dir=".git"):
 350     """
 351     Recursively determines the disk usage of a directory, excluding
 352     .git directories.  Value is in bytes.  If ``dir`` is omitted, the
 353     current working directory is assumed.
 354     """
 355     if dir is None: dir = os.getcwd()
 356     sum_sizes = 0
 357     for root, _, files in os.walk(dir):
 358         for name in files:
 359             if not os.path.join(root, name).startswith(os.path.join(dir, excluded_dir)):
 360                 file = os.path.join(root, name)
 361                 try:
 362                     if os.path.islink(file): continue
 363                     sum_sizes += os.path.getsize(file)
 364                 except OSError as e:
 365                     if e.errno == errno.ENOENT:
 366                         logging.warning("%s disappeared before we could stat", file)
 367                     else:
 368                         raise
 369     return sum_sizes
 370
 371 def random_key(length=30):
 372     """Generates a random alphanumeric key of ``length`` size."""
 373     return ''.join(random.choice(string.letters + string.digits) for i in xrange(length))
 374
 375 def truncate(version):
 376     """Truncates the Scripts specific version number."""
 377     return str(version).partition('-scripts')[0]
 378
 379 def init_wizard_dir():
 380     """
 381     Generates a .wizard directory and initializes it with some common
 382     files.  This operation is idempotent.
 383     """
 384     # no harm in doing this repeatedly
 385     wizard_dir = ".wizard"
 386     if not os.path.isdir(wizard_dir):
 387         os.mkdir(wizard_dir)
 388     open(os.path.join(wizard_dir, ".htaccess"), "w").write("Deny from all\n")
 389     open(os.path.join(wizard_dir, ".gitignore"), "w").write("*\n")
 390
 391 class NoOperatorInfo(wizard.Error):
 392     """No information could be found about the operator from Kerberos."""
 393     pass
 394
 395 class PermissionsError(IOError):
 396     errno = errno.EACCES
 397
 398 class NoSuchDirectoryError(IOError):
 399     errno = errno.ENOENT
 400
 401 class DirectoryLockedError(wizard.Error):
 402     def __init__(self, dir):
 403         self.dir = dir
 404     def __str__(self):
 405         return """
 406
 407 ERROR: Could not acquire lock on directory.  Maybe there is
 408 another migration process running?
 409 """
 410
 411 class DNSError(socket.gaierror):
 412     errno = socket.EAI_NONAME
 413     #: Hostname that could not resolve name
 414     host = None
 415     def __init__(self, host):
 416         self.host = host
 417     def __str__(self):
 418         return """
 419
 420 ERROR: Could not resolve hostname %s.
 421 """ % self.host