2 Miscellaneous utility functions and classes.
6 from wizard.util import *
27 class ChangeDirectory(object):
29 Context for temporarily changing the working directory.
31 >>> with ChangeDirectory("/tmp"):
35 def __init__(self, dir):
39 self.olddir = os.getcwd()
41 def __exit__(self, *args):
44 class Counter(object):
46 Object for counting different values when you don't know what
47 they are a priori. Supports index access and iteration.
49 >>> counter = Counter()
50 >>> counter.count("foo")
51 >>> print counter["foo"]
56 def count(self, value):
57 """Increments count for ``value``."""
58 self.dict.setdefault(value, 0)
60 def __getitem__(self, key):
63 return self.dict.__iter__()
65 """Returns the max counter value seen."""
66 return max(self.dict.values())
68 """Returns the sum of all counter values."""
69 return sum(self.dict.values())
71 """Returns the keys of counters."""
72 return self.dict.keys()
74 class PipeToLess(object):
76 Context for printing output to a pager. Use this if output
77 is expected to be long.
80 self.proc = subprocess.Popen("less", stdin=subprocess.PIPE)
81 self.old_stdout = sys.stdout
82 sys.stdout = self.proc.stdin
83 def __exit__(self, *args):
85 self.proc.stdin.close()
87 sys.stdout = self.old_stdout
89 class IgnoreKeyboardInterrupts(object):
91 Context for temporarily ignoring keyboard interrupts. Use this
92 if aborting would cause more harm than finishing the job.
95 signal.signal(signal.SIGINT,signal.SIG_IGN)
96 def __exit__(self, *args):
97 signal.signal(signal.SIGINT, signal.default_int_handler)
99 class LockDirectory(object):
101 Context for locking a directory.
103 def __init__(self, lockfile, expiry = 3600):
104 self.lockfile = lockfile
105 self.expiry = expiry # by default an hour
108 for i in range(0, 3):
110 os.open(self.lockfile, os.O_CREAT | os.O_EXCL)
111 open(self.lockfile, "w").write("%d" % os.getpid())
113 if e.errno == errno.EEXIST:
114 # There is a possibility of infinite recursion, but we
115 # expect it to be unlikely, and not harmful if it does happen
116 with LockDirectory(self.lockfile + "_"):
117 # See if we can break the lock
119 pid = open(self.lockfile, "r").read().strip()
120 if not os.path.exists("/proc/%s" % pid):
121 # break the lock, try again
122 logging.warning("Breaking orphaned lock at %s", self.lockfile)
123 os.unlink(self.lockfile)
126 # check if the file is expiry old, if so, break the lock, try again
127 if time.time() - os.stat(self.lockfile).st_mtime > self.expiry:
128 logging.warning("Breaking stale lock at %s", self.lockfile)
129 os.unlink(self.lockfile)
132 if e.errno == errno.ENOENT:
136 # oh hey, it went away; try again
138 raise DirectoryLockedError(os.getcwd())
139 elif e.errno == errno.EACCES:
140 raise PermissionsError(os.getcwd())
143 raise DirectoryLockedError(os.getcwd())
144 def __exit__(self, *args):
146 os.unlink(self.lockfile)
152 Changes a directory, but has special exceptions for certain
158 if e.errno == errno.EACCES:
159 raise PermissionsError()
160 elif e.errno == errno.ENOENT:
161 raise NoSuchDirectoryError()
166 A map function for dictionaries. Only changes values.
168 >>> dictmap(lambda x: x + 2, {'a': 1, 'b': 2})
171 return dict((k,f(v)) for k,v in d.items())
175 A map function for dictionaries that passes key and value.
177 >>> dictkmap(lambda x, y: x + y, {1: 4, 3: 4})
180 return dict((k,f(k,v)) for k,v in d.items())
182 def get_exception_name(output):
184 Reads the traceback from a Python program and grabs the
185 fully qualified exception name.
187 lines = output.split("\n")
190 for line in lines[1:]:
192 if not line: continue
198 return line.partition(':')[0]
201 def get_dir_uid(dir):
202 """Finds the uid of the person who owns this directory."""
203 return os.stat(dir).st_uid
205 def get_dir_owner(dir = "."):
207 Finds the name of the locker this directory is in.
211 This function uses the passwd database and thus
212 only works on scripts servers when querying directories
215 uid = get_dir_uid(dir)
217 pwentry = pwd.getpwuid(uid)
218 return pwentry.pw_name
220 # do an pts query to get the name
221 return subprocess.Popen(['pts', 'examine', str(uid)], stdout=subprocess.PIPE).communicate()[0].partition(",")[0].partition(": ")[2]
224 """Returns the commit ID of the current Wizard install."""
225 # If you decide to convert this to use wizard.shell, be warned
226 # that there is a circular dependency, so this function would
227 # probably have to live somewhere else, probably wizard.git
228 wizard_git = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".git")
229 return subprocess.Popen(["git", "--git-dir=" + wizard_git, "rev-parse", "HEAD"], stdout=subprocess.PIPE).communicate()[0].rstrip()
231 def get_operator_info():
233 Returns tuple of ``(realname, email)`` about the person running
234 the script. If run from a scripts server, get info from Hesiod.
235 Otherwise, use the passwd database (email generated probably won't
236 actually accept mail). Useful when generating commit messages.
238 username = get_operator_name_from_gssapi()
241 hesinfo = subprocess.Popen(["hesinfo", username, "passwd"],stdout=subprocess.PIPE).communicate()[0]
242 fields = hesinfo.partition(",")[0]
243 realname = fields.rpartition(":")[2]
244 return realname, username + "@mit.edu"
246 # more traditional approach, but the email probably doesn't work
249 # since root isn't actually a useful designation, but maybe
250 # SUDO_USER contains something helpful
251 sudo_user = os.getenv("SUDO_USER")
254 pwdentry = pwd.getpwnam(sudo_user)
256 pwdentry = pwd.getpwuid(uid)
257 # XXX: error checking might be nice
258 # We follow the Ubuntu convention of gecos being a comma split field
259 # with the person's realname being the first entry.
260 return pwdentry.pw_gecos.split(",")[0], pwdentry.pw_name + "@" + socket.gethostname()
262 def get_operator_git():
264 Returns ``Real Name <username@mit.edu>`` suitable for use in
265 Git ``Something-by:`` string.
267 return "%s <%s>" % get_operator_info()
269 def get_operator_name_from_gssapi():
271 Returns username of the person operating this script based
272 off of the :envvar:`SSH_GSSAPI_NAME` environment variable.
276 :envvar:`SSH_GSSAPI_NAME` is not set by a vanilla OpenSSH
277 distributions. Scripts servers are patched to support this
278 environment variable.
280 principal = os.getenv("SSH_GSSAPI_NAME")
283 instance, _, _ = principal.partition("@")
284 if instance.endswith("/root"):
285 username, _, _ = principal.partition("/")
290 def set_operator_env():
292 Sets :envvar:`GIT_COMMITTER_NAME` and :envvar:`GIT_COMMITTER_EMAIL`
293 environment variables if applicable. Does nothing if
294 :func:`get_operator_info` throws :exc:`NoOperatorInfo`.
297 op_realname, op_email = get_operator_info()
298 os.putenv("GIT_COMMITTER_NAME", op_realname)
299 os.putenv("GIT_COMMITTER_EMAIL", op_email)
300 except NoOperatorInfo:
303 def set_author_env():
305 Sets :envvar:`GIT_AUTHOR_NAME` and :envvar:`GIT_AUTHOR_EMAIL` environment
306 variables if applicable. Does nothing if :func:`get_dir_owner` fails.
309 # XXX: should check if the directory is in AFS, and if not, use
310 # a more traditional metric
311 lockername = get_dir_owner()
312 os.putenv("GIT_AUTHOR_NAME", "%s locker" % lockername)
313 os.putenv("GIT_AUTHOR_EMAIL", "%s@scripts.mit.edu" % lockername)
314 except KeyError: # XXX: This doesn't actually make sense
318 """Sets all appropriate environment variables for Git commits."""
322 def get_git_footer():
323 """Returns strings for placing in Git log info about Wizard."""
324 return "\n".join(["Wizard-revision: %s" % get_revision()
325 ,"Wizard-args: %s" % " ".join(sys.argv)
328 def safe_unlink(file):
329 """Moves a file/dir to a backup location."""
330 if not os.path.exists(file):
332 prefix = "%s.bak" % file
334 for i in itertools.count():
335 name = "%s.%d" % (prefix, i)
336 if not os.path.exists(name):
338 os.rename(file, name)
341 def soft_unlink(file):
342 """Unlink a file, but don't complain if it doesn't exist."""
348 def fetch(host, path, subpath, post=None):
350 # XXX: Special case if it's https; not sure why this data isn't
352 h = httplib.HTTPConnection(host)
353 fullpath = path.rstrip("/") + "/" + subpath.lstrip("/") # to be lenient about input we accept
355 headers = {"Content-type": "application/x-www-form-urlencoded"}
356 h.request("POST", fullpath, urllib.urlencode(post), headers)
358 h.request("GET", fullpath)
363 except socket.gaierror as e:
364 if e.errno == socket.EAI_NONAME:
369 def mixed_newlines(filename):
370 """Returns ``True`` if ``filename`` has mixed newlines."""
371 f = open(filename, "U") # requires universal newline support
373 ret = isinstance(f.newlines, tuple)
374 f.close() # just to be safe
377 def random_key(length=30):
378 """Generates a random alphanumeric key of ``length`` size."""
379 return ''.join(random.choice(string.letters + string.digits) for i in xrange(length))
381 class NoOperatorInfo(wizard.Error):
382 """No information could be found about the operator from Kerberos."""
385 class PermissionsError(IOError):
388 class NoSuchDirectoryError(IOError):
391 class DirectoryLockedError(wizard.Error):
392 def __init__(self, dir):
397 ERROR: Could not acquire lock on directory. Maybe there is
398 another migration process running?
401 class DNSError(socket.gaierror):
402 errno = socket.EAI_NONAME
403 #: Hostname that could not resolve name
405 def __init__(self, host):
410 ERROR: Could not resolve hostname %s.