More bugfixes from running live.

[wizard.git] / wizard / scripts.py

diff --git a/wizard/scripts.py b/wizard/scripts.py
index 35b9ba942db0905aa5ab8e37ac03b9188cde4c97..a81f539ed641a4810a025192217dc6c5146ba72c 100644 (file)
--- a/wizard/scripts.py
+++ b/wizard/scripts.py
@@ -1,8 +1,10 @@
 import os
+import shlex
 
-from wizard import shell
+import wizard
+from wizard import shell, util
 
-def get_sql_credentials():
+def get_sql_credentials(vars=None):
     """
     Attempts to determine a user's MySQL credentials.  They are
     returned as a three-tuple (host, user, password).
@@ -13,8 +15,14 @@ def get_sql_credentials():
     password = os.getenv("WIZARD_MYSQL_PASSWORD")
     if host is not None and user is not None and password is not None:
         return (host, user, password)
+    # XXX: this is very fragile
+    elif vars and "WIZARD_DBSERVER" in vars and "WIZARD_DBUSER" in vars and "WIZARD_DBPASSWORD" in vars:
+        return (shlex.split(vars[x])[0] for x in ("WIZARD_DBSERVER", "WIZARD_DBUSER", "WIZARD_DBPASSWORD"))
     try:
-        return sh.eval("/mit/scripts/sql/bin/get-password").split()
+        tuple = sh.eval("/mit/scripts/sql/bin/get-password").split()
+        if len(tuple) == 3:
+            return tuple
+        return None
     except CallError:
         return None
 
@@ -36,3 +44,55 @@ def get_web_host_and_path(dir=None):
         return None
     return (util.get_dir_owner(dir) + ".scripts.mit.edu", web_path)
 
+def get_quota_usage_and_limit(dir=None):
+    """
+    Returns a tuple (quota usage, quota limit).  Works only for scripts
+    servers.  Values are in KiB.  Returns ``(0, None)`` if we couldn't figure it out.
+    """
+    # XXX: The correct way is to implement Python modules implementing
+    # bindings for all the appropriate interfaces
+    def parse_last_quote(ret):
+        return ret.rstrip('\'').rpartition('\'')[2]
+    if dir is None:
+        dir = os.getcwd()
+    sh = shell.Shell()
+    try:
+        cell = parse_last_quote(sh.eval("fs", "whichcell", "-path", dir))
+    except shell.CallError:
+        return (0, None)
+    mount = None
+    while dir:
+        try:
+            volume = parse_last_quote(sh.eval("fs", "lsmount", dir))[1:]
+            break
+        except shell.CallError:
+            dir = os.path.dirname(dir)
+    if not volume: return (0, None)
+    try:
+        result = sh.eval("vos", "examine", "-id", volume, "-cell", cell).splitlines()
+    except shell.CallError:
+        return (0, None)
+    try:
+        usage = int(result[0].split()[3])
+        limit = int(result[3].split()[1]) # XXX: FRAGILE
+    except ValueError:
+        raise QuotaParseError("vos examine output was:\n\n" + "\n".join(result))
+    return (usage, limit)
+
+# XXX: Possibly in the wrong module
+def get_disk_usage(dir=None):
+    """
+    Recursively determines the disk usage of a directory, excluding
+    .git directories.  Value is in bytes.
+    """
+    if dir is None: dir = os.getcwd()
+    return int(shell.Shell().eval("du", "--exclude=.git", "-s", dir).split()[0])
+
+class QuotaParseError(wizard.Error):
+    def __init__(self, msg):
+        self.msg = msg
+    def __str__(self):
+        return """
+
+ERROR: Could not parse quota. %s
+""" % self.msg