Minor doc updates and code refinements.

[wizard.git] / wizard / scripts.py

diff --git a/wizard/scripts.py b/wizard/scripts.py
index 415be1a064a27bec416d015a0dd1eb56e5f7b4bb..66c62d887782dd01f0feb973b1c0154931602487 100644 (file)
--- a/wizard/scripts.py
+++ b/wizard/scripts.py
@@ -1,13 +1,113 @@
-from wizard import shell
+import os
+import shlex
+import errno
+import logging
 
-def get_sql_credentials():
+import wizard
+from wizard import shell, util
+
+def get_sql_credentials(vars=None):
     """
     Attempts to determine a user's MySQL credentials.  They are
     returned as a three-tuple (host, user, password).
     """
     sh = shell.Shell()
+    host = os.getenv("WIZARD_MYSQL_HOST")
+    user = os.getenv("WIZARD_MYSQL_USER")
+    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()
-    except CallError:
+        tuple = sh.eval("/mit/scripts/sql/bin/get-password").split()
+        if len(tuple) == 3:
+            return tuple
+        return None
+    except shell.CallError:
+        return None
+
+def get_web_host_and_path(dir=None):
+    """
+    Attempts to determine webhost and path for the current directory
+    as it would be accessible from the web.  Works only for scripts
+    servers.  Returns a tuple web_host, web_path, or None if it failed.
+    """
+    # XXX: THIS CODE SUCKS
+    host = os.getenv("WIZARD_WEB_HOST")
+    path = os.getenv("WIZARD_WEB_PATH")
+    if host is not None and path is not None:
+        return (host, path.rstrip('/'))
+    if not dir:
+        dir = os.getcwd()
+    # XXX: possible security risk?
+    homedir, _, web_path = dir.partition("/web_scripts")
+    if not web_path:
         return None
+    return (util.get_dir_owner(homedir) + ".scripts.mit.edu", web_path.rstrip('/'))
+
+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)
+    except OSError as e:
+        if e.errno == errno.ENOENT:
+            return (0, None)
+        raise
+    mount = None
+    while dir:
+        try:
+            volume = parse_last_quote(sh.eval("fs", "lsmount", dir))[1:]
+            break
+        except shell.CallError:
+            dir = os.path.dirname(dir)
+        except OSError as e:
+            if e.errno == errno.ENOENT:
+                return (0, None)
+            raise
+    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, excluded_dir=".git"):
+    """
+    Recursively determines the disk usage of a directory, excluding
+    .git directories.  Value is in bytes.
+    """
+    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(dir + excluded_dir):
+                sum_sizes += os.path.getsize(os.path.join(root, name))
+    return sum_sizes
+class QuotaParseError(wizard.Error):
+    def __init__(self, msg):
+        self.msg = msg
+    def __str__(self):
+        return """
 
+ERROR: Could not parse quota. %s
+""" % self.msg