#!/usr/bin/env python3

import socket
import sys
import os
import re
import argparse
import tempfile


# Open a lock so this script can't run when it's already running. This method only works on Linux.
def get_lock(process_name):
    # Without holding a reference to our socket somewhere it gets garbage
    # collected when the function exits
    get_lock._lock_socket = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)

    try:
        # The null byte (\0) means the the socket is created
        # in the abstract namespace instead of being created
        # on the file system itself.
        # Works only in Linux
        get_lock._lock_socket.bind('\0' + process_name)
    except socket.error:
        print("Lock exists - %s is already running" % process_name, file=sys.stderr)
        sys.exit(1)


# Get the usage and quota for a given maildir. Returns a tuple (used_bytes, quota_bytes, is_over_quota)
# Quota can legitimately be zero (no quota set - e.g. unlimited, or no mail received yet from smtp in), so this is not overquota
def get_quota(maildir):
    quotafile = os.path.join(maildir, 'maildirsize')
    if not os.path.isfile(quotafile):
        # If they don't have a maildirsize file, they haven't received any mail, so they can't be overquota
        return (0, 0, False)

    f = open(quotafile, 'r')
    match = re.search('^(\d+)S', f.readline())
    if not match:
        raise ValueError("Quota file %s does not begin with quota bytes" % quotafile)

    quota_bytes = int(match[1])

    # bail out early here if they don't have a quota - we don't need to count how many bytes they are using (unless we're verbose)
    if quota_bytes == 0 and not verbose:
        return (0, 0, False)

    used_bytes = 0
    for line in f:
        match = re.search('^(-?\d+)\s', line)
        if match:
            used_bytes += int(match[1])

    return (used_bytes, quota_bytes, quota_bytes > 0 and (used_bytes >= quota_bytes))


# Find all users who are over quota and print to the given file
def find_over_quota(output_file):

    for domain in os.listdir(vmail_dir):
        if domain == 'lost+found': continue

        domain_dir = os.path.join(vmail_dir, domain)
        if os.path.isdir(domain_dir):

            for localpart in os.listdir(domain_dir):
                user_dir = os.path.join(domain_dir, localpart)
                if os.path.isdir(user_dir):
                    address = localpart + '@' + domain

                    maildir_dir = os.path.join(user_dir, 'Maildir')
                    if os.path.isdir(maildir_dir):
                        (used_bytes, quota_bytes, over_quota) = get_quota(maildir_dir)
                        if verbose:
                            print("# %s using %d of %d (%s)" % (address, used_bytes, quota_bytes, over_quota))
                        if over_quota:
                            print(address, file=output_file)



parser = argparse.ArgumentParser(description='Check for virtual users that are overquota.')

parser.add_argument('-v', '--verbose',
                    default=False,
                    action='store_const', const=True,
                    help='Print maildir size and quota for each user to stdout.')

parser.add_argument('-o', '--output-file', nargs=1,
                    help='File to output to. A temporary file is created and renamed atomically into place.')

parser.add_argument('-d', '--dir', nargs=1, required=True,
                    help='The directory containing the domains which themselves contain users.')


args = parser.parse_args()

output_filename = args.output_file[0] if args.output_file else None
vmail_dir = args.dir[0]
verbose = args.verbose

get_lock(os.path.basename(__file__))


if output_filename:
    temp_file = tempfile.NamedTemporaryFile('w',
                                            dir=os.path.dirname(output_filename),
                                            prefix=os.path.basename(output_filename) + '_')

    find_over_quota(temp_file)

    # atomically move the temporary file into place
    os.chmod(temp_file.name, 0o644)
    os.replace(temp_file.name, output_filename)

    # recreate the temporary file so that the temporaryfile closer can happily delete it!
    os.mknod(temp_file.name)
else:
    find_over_quota(sys.stdout)