X-Git-Url: https://scripts.mit.edu/gitweb/wizard.git/blobdiff_plain/b29d182f610f68dcdc9525e08d5ad4fab65ffaa4..c1b068b03a604652f3eb284fdb92bcb1171d9601:/wizard/util.py diff --git a/wizard/util.py b/wizard/util.py index eca883e..c7d8969 100644 --- a/wizard/util.py +++ b/wizard/util.py @@ -17,6 +17,10 @@ import itertools import signal import httplib import urllib +import time +import logging +import random +import string import wizard @@ -57,6 +61,15 @@ class Counter(object): return self.dict[key] def __iter__(self): return self.dict.__iter__() + def max(self): + """Returns the max counter value seen.""" + return max(self.dict.values()) + def sum(self): + """Returns the sum of all counter values.""" + return sum(self.dict.values()) + def keys(self): + """Returns the keys of counters.""" + return self.dict.keys() class PipeToLess(object): """ @@ -87,17 +100,47 @@ class LockDirectory(object): """ Context for locking a directory. """ - def __init__(self, lockfile): + def __init__(self, lockfile, expiry = 3600): self.lockfile = lockfile + self.expiry = expiry # by default an hour def __enter__(self): - try: - os.open(self.lockfile, os.O_CREAT | os.O_EXCL) - except OSError as e: - if e.errno == errno.EEXIST: - raise DirectoryLockedError(os.getcwd()) - elif e.errno == errno.EACCES: - raise PermissionsError(os.getcwd()) - raise + # It's A WAVY + for i in range(0, 3): + try: + os.open(self.lockfile, os.O_CREAT | os.O_EXCL) + open(self.lockfile, "w").write("%d" % os.getpid()) + except OSError as e: + if e.errno == errno.EEXIST: + # There is a possibility of infinite recursion, but we + # expect it to be unlikely, and not harmful if it does happen + with LockDirectory(self.lockfile + "_"): + # See if we can break the lock + try: + pid = open(self.lockfile, "r").read().strip() + if not os.path.exists("/proc/%s" % pid): + # break the lock, try again + logging.warning("Breaking orphaned lock at %s", self.lockfile) + os.unlink(self.lockfile) + continue + try: + # check if the file is expiry old, if so, break the lock, try again + if time.time() - os.stat(self.lockfile).st_mtime > self.expiry: + logging.warning("Breaking stale lock at %s", self.lockfile) + os.unlink(self.lockfile) + continue + except OSError as e: + if e.errno == errno.ENOENT: + continue + raise + except IOError: + # oh hey, it went away; try again + continue + raise DirectoryLockedError(os.getcwd()) + elif e.errno == errno.EACCES: + raise PermissionsError(os.getcwd()) + raise + return + raise DirectoryLockedError(os.getcwd()) def __exit__(self, *args): try: os.unlink(self.lockfile) @@ -165,13 +208,17 @@ def get_dir_owner(dir = "."): .. note:: - This function uses the passwd database and thus - only works on scripts servers when querying directories - that live on AFS. + When querying AFS servers, this function only works if + you're on a Scripts server (which has the correct passwd + database) or if you're on a Debathena machine. """ - pwentry = pwd.getpwuid(get_dir_uid(dir)) - # XXX: Error handling! - return pwentry.pw_name + uid = get_dir_uid(dir) + try: + pwentry = pwd.getpwuid(uid) + return pwentry.pw_name + except KeyError: + # do an pts query to get the name + return subprocess.Popen(['pts', 'examine', str(uid)], stdout=subprocess.PIPE).communicate()[0].partition(",")[0].partition(": ")[2] def get_revision(): """Returns the commit ID of the current Wizard install.""" @@ -280,29 +327,91 @@ def get_git_footer(): def safe_unlink(file): """Moves a file/dir to a backup location.""" - if not os.path.exists(file): + if not os.path.lexists(file): return None prefix = "%s.bak" % file name = None for i in itertools.count(): name = "%s.%d" % (prefix, i) - if not os.path.exists(name): + if not os.path.lexists(name): break os.rename(file, name) return name +def soft_unlink(file): + """Unlink a file, but don't complain if it doesn't exist.""" + try: + os.unlink(file) + except OSError: + pass + +def makedirs(path): + """ + Create a directory path (a la ``mkdir -p`` or ``os.makedirs``), + but don't complain if it already exists. + """ + try: + os.makedirs(path) + except OSError as exc: + if exc.errno == errno.EEXIST: + pass + else: + raise + def fetch(host, path, subpath, post=None): - h = httplib.HTTPConnection(host) - fullpath = path + "/" + subpath - if post: - headers = {"Content-type": "application/x-www-form-urlencoded"} - h.request("POST", fullpath, urllib.urlencode(post), headers) - else: - h.request("GET", fullpath) - r = h.getresponse() - data = r.read() - h.close() - return data + try: + # XXX: Special case if it's https; not sure why this data isn't + # passed + h = httplib.HTTPConnection(host) + fullpath = path.rstrip("/") + "/" + subpath.lstrip("/") # to be lenient about input we accept + if post: + headers = {"Content-type": "application/x-www-form-urlencoded"} + h.request("POST", fullpath, urllib.urlencode(post), headers) + else: + h.request("GET", fullpath) + r = h.getresponse() + data = r.read() + h.close() + return data + except socket.gaierror as e: + if e.errno == socket.EAI_NONAME: + raise DNSError(host) + else: + raise + +def mixed_newlines(filename): + """Returns ``True`` if ``filename`` has mixed newlines.""" + f = open(filename, "U") # requires universal newline support + f.read() + ret = isinstance(f.newlines, tuple) + f.close() # just to be safe + return ret + +def disk_usage(dir=None, excluded_dir=".git"): + """ + Recursively determines the disk usage of a directory, excluding + .git directories. Value is in bytes. If ``dir`` is omitted, the + current working directory is assumed. + """ + if dir is None: dir = os.getcwd() + sum_sizes = 0 + for root, _, files in os.walk(dir): + for name in files: + if not os.path.join(root, name).startswith(os.path.join(dir, excluded_dir)): + file = os.path.join(root, name) + try: + if os.path.islink(file): continue + sum_sizes += os.path.getsize(file) + except OSError as e: + if e.errno == errno.ENOENT: + logging.warning("%s disappeared before we could stat", file) + else: + raise + return sum_sizes + +def random_key(length=30): + """Generates a random alphanumeric key of ``length`` size.""" + return ''.join(random.choice(string.letters + string.digits) for i in xrange(length)) class NoOperatorInfo(wizard.Error): """No information could be found about the operator from Kerberos.""" @@ -324,3 +433,14 @@ ERROR: Could not acquire lock on directory. Maybe there is another migration process running? """ +class DNSError(socket.gaierror): + errno = socket.EAI_NONAME + #: Hostname that could not resolve name + host = None + def __init__(self, host): + self.host = host + def __str__(self): + return """ + +ERROR: Could not resolve hostname %s. +""" % self.host