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         logging.debug("Response code: %d", r.status)
 333         logging.debug("Response headers: %s", r.msg)
 334         data = r.read()
 335         h.close()
 336         return data
 337     except socket.gaierror as e:
 338         if e.errno == socket.EAI_NONAME:
 339             raise DNSError(host)
 340         else:
 341             raise
 342
 343 def mixed_newlines(filename):
 344     """Returns ``True`` if ``filename`` has mixed newlines."""
 345     f = open(filename, "U") # requires universal newline support
 346     f.read()
 347     ret = isinstance(f.newlines, tuple)
 348     f.close() # just to be safe
 349     return ret
 350
 351 def disk_usage(dir=None, excluded_dir=".git"):
 352     """
 353     Recursively determines the disk usage of a directory, excluding
 354     .git directories.  Value is in bytes.  If ``dir`` is omitted, the
 355     current working directory is assumed.
 356     """
 357     if dir is None: dir = os.getcwd()
 358     sum_sizes = 0
 359     for root, _, files in os.walk(dir):
 360         for name in files:
 361             if not os.path.join(root, name).startswith(os.path.join(dir, excluded_dir)):
 362                 file = os.path.join(root, name)
 363                 try:
 364                     if os.path.islink(file): continue
 365                     sum_sizes += os.path.getsize(file)
 366                 except OSError as e:
 367                     if e.errno == errno.ENOENT:
 368                         logging.warning("%s disappeared before we could stat", file)
 369                     else:
 370                         raise
 371     return sum_sizes
 372
 373 def random_key(length=30):
 374     """Generates a random alphanumeric key of ``length`` size."""
 375     return ''.join(random.choice(string.letters + string.digits) for i in xrange(length))
 376
 377 def truncate(version):
 378     """Truncates the Scripts specific version number."""
 379     return str(version).partition('-scripts')[0]
 380
 381 def init_wizard_dir():
 382     """
 383     Generates a .wizard directory and initializes it with some common
 384     files.  This operation is idempotent.
 385     """
 386     # no harm in doing this repeatedly
 387     wizard_dir = ".wizard"
 388     if not os.path.isdir(wizard_dir):
 389         os.mkdir(wizard_dir)
 390     open(os.path.join(wizard_dir, ".htaccess"), "w").write("Deny from all\n")
 391     open(os.path.join(wizard_dir, ".gitignore"), "w").write("*\n")
 392
 393 class NoOperatorInfo(wizard.Error):
 394     """No information could be found about the operator from Kerberos."""
 395     pass
 396
 397 class PermissionsError(IOError):
 398     errno = errno.EACCES
 399
 400 class NoSuchDirectoryError(IOError):
 401     errno = errno.ENOENT
 402
 403 class DirectoryLockedError(wizard.Error):
 404     def __init__(self, dir):
 405         self.dir = dir
 406     def __str__(self):
 407         return """
 408
 409 ERROR: Could not acquire lock on directory.  Maybe there is
 410 another migration process running?
 411 """
 412
 413 class DNSError(socket.gaierror):
 414     errno = socket.EAI_NONAME
 415     #: Hostname that could not resolve name
 416     host = None
 417     def __init__(self, host):
 418         self.host = host
 419     def __str__(self):
 420         return """
 421
 422 ERROR: Could not resolve hostname %s.
 423 """ % self.host