afs-quota-notify

#!/usr/bin/python3


import os
import re
import pwd
import sys
import getopt
import socket
import smtplib
import subprocess
import configparser
from io import StringIO

from email.mime.text import MIMEText

SCRIPT_NAME = 'afs-quota-notify'
SECTION = 'afs-quota-notify'

USAGE = r'''
Usage: {SCRIPT_NAME} [options] [<config file>]

Send email if AFS directories are close to being over quota. The values
'email', 'threshold', and 'directories' are read from the '{SECTION}' section
of an ini-style config file, or can be passed via command-line options.

Options:

    -d <directories>      A comma or space separated list of the directories
                          to check

    -e <email addresses>  A comma or space separated list of who to email

    -h                    Print this message

    -n                    Dry run: don't actually email, just print what would
                          be sent

    -t <percent>          The warning threshold as a percentage from 0.0% to
                          100.0%
'''.format(**globals())

DEFAULT_CONFIG = r'''
[{SECTION}]

email = matyas+cron@cs.wisc.edu

directories = /p/vdt/public/html
              /p/vdt/public/html/tarball-client
              /p/vdt/public/html/upstream
              /p/condor/workspaces/vdt
              /p/condor/workspaces/vdt/git

threshold = 90.0
'''.format(**globals())


def send_email(recipients, subject, text, dry_run=False):
    msg = MIMEText(text)

    username = pwd.getpwuid(os.getuid()).pw_name
    hostname = socket.gethostname()

    msg['Subject'] = subject
    msg['To'] = ', '.join(recipients)
    msg['From'] = username + "@" + hostname

    if dry_run:
        print(msg.as_string())
    else:
        smtp = smtplib.SMTP('localhost')
        smtp.sendmail(msg['From'], recipients, msg.as_string())
        smtp.quit()


def _disk_quota_percent_used(listquota_line):
    """Helper for decoding the output of 'fs listquota'. Takes a line such as:
    u.matyas-497381             20000000   12209872   61%          0%
    and returns the percentage of quota used (as a float, out of 100.0).  (Does
    not use the percentage from the line, but calculates it from the total and
    used (the first two numbers)).  Can return TypeError or IndexError if the
    line is not in the expected format.

    """
    line_split = re.split(r'\s+', listquota_line)
    quota = line_split[1]
    used = line_split[2]
    percent = (float(used) / float(quota)) * 100.0

    return percent


def get_disk_usage_percent(directory):
    """Return the percent of disk quota used in the directory. 'directory'
    must be on AFS. Uses 'fs listquota' to get the quota remaining.

    """
    listquota_proc = subprocess.Popen(
        ['/usr/bin/fs', 'listquota', directory], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    listquota_output = listquota_proc.communicate()[0].decode("latin-1").split("\n")
    if listquota_proc.returncode != 0:
        raise RuntimeError('{0}: fs listquota returned {1:d}. Output: {2}'.format(directory, listquota_proc.returncode, listquota_output))

    # The second line contains the information
    try:
        return _disk_quota_percent_used(listquota_output[1])
    except (TypeError, IndexError, ValueError):
        raise RuntimeError('{0}: fs listquota did not return expected output: {1}'.format(directory, listquota_output))


def read_config(config_fp):
    config = configparser.ConfigParser()
    config.read_file(config_fp)

    directories = re.split(r'[ ,\n]+', config.get(SECTION, 'directories'))
    email = re.split(r'[ ,\n]+', config.get(SECTION, 'email'))
    threshold = config.getfloat(SECTION, 'threshold')

    return directories, email, threshold


def main(argv):
    dry_run = False

    try:
        optlist, args = getopt.gnu_getopt(argv[1:], 'd:e:hnt:')
    except getopt.error as err:
        print(str(err))
        print(USAGE)
        return 2

    config_text = "Using default config"
    if len(args) > 0:
        filename = os.path.realpath(args[0])
        if filename.startswith("/afs/cs.wisc.edu"):
            filename = filename[16:]
        config_fp = open(filename, 'rt')
        config_text = "Using config at " + filename
    else:
        config_fp = StringIO(DEFAULT_CONFIG)
    directories, email, threshold = read_config(config_fp)
    config_fp.close()

    for opt, optarg in optlist:
        if opt == "-d":
            directories = re.split(r'[ ,]+', optarg)
        elif opt == "-e":
            email = re.split(r'[ ,]+', optarg)
        elif opt == "-h":
            print(USAGE)
            return 0
        elif opt == "-n":
            dry_run = True
        elif opt == "-t":
            threshold = float(optarg)

    if not directories:
        print("No directories given!", file=sys.stderr)
        print(USAGE)
        return 2

    dir_column_width = max([len(d) for d in directories])

    msg_text = config_text + "\n\n"
    msg_subject = "AFS quota monitor: "

    warning_count = 0
    highest = 0.0
    highest_path = ""
    errors = ""
    for dirpath in directories:
        try:
            disk_usage = get_disk_usage_percent(dirpath)
        except RuntimeError as err:
            errors += "%s\n" % err
            continue

        warning_str = ""

        if disk_usage > highest:
            highest, highest_path = disk_usage, dirpath

        if disk_usage > threshold:
            warning_str = " << ABOVE THRESHOLD"
            warning_count += 1

        # Body text
        # Examples:
        # /p/condor/vdt/workspaces  50.0%
        # /u/m/a/matyas             71.0% << ABOVE THRESHOLD
        msg_text += "{0:<{width}} {1:>5.1f}%{2}\n".format(
            dirpath, disk_usage, warning_str, width=dir_column_width)

    # Subject
    # Examples:
    # /u/m/a/matyas and 1 other dir are above 40.0%
    # /u/m/a/matyas is above 70.0%
    # all dirs are below 90.0%
    if warning_count > 1:
        other_dirs = warning_count - 1
        msg_subject += "{0} and {1:d} other{3} dir{3} are above {2:.1f}%".format(
            highest_path,
            other_dirs,
            threshold,
            's' if other_dirs > 1 else '')

    elif warning_count == 1:
        msg_subject += "{0} is above {1:.1f}%".format(highest_path, threshold)

    else:
        msg_subject += "all dirs are below {0:.1f}%".format(threshold)

    if errors:
        msg_subject += "; ERRORS ENCOUNTERED"
        msg_text += "\n\nThe following errors were encountered:\n\n" + str(errors)

    send_email(email, msg_subject, msg_text, dry_run)


if __name__ == '__main__':
    sys.exit(main(sys.argv))