#!/usr/bin/python
#
# This script is released to the public domain.
# Example usage:
# $ imapquota.py -s imap.mail.drexel.edu xyz12@drexel.edu
# Password:
# Used 92693 of 102400 KB
# 90.5 percent

import getpass
import imaplib
import re


def parse_server(server, ssl=False):
    server_bits = server.split(':')
    assert len(server_bits) in [1,2], 'invalid server: %s' % server_bits
    server = server_bits[0]
    if len(server_bits) == 2:
        port = server_bits[1]
    elif ssl == True:
        port = 993
    else:
        port = 143
    return (server, port)

def traverse_mailboxes(imap_connection, root='""', recursive=False):
    status,value = imap_connection.list(root, '%')
    regexp = re.compile(r'\((.*)\) "(.)" (.*)')
    #status,value = ('OK', [
    #    '(\NoInferiors) "/" INBOX',
    #    '(\HasNoChildren) "/" APS',
    #    ...])
    for mailbox_line in value:
        if mailbox_line is None:
            continue
        match = regexp.match(mailbox_line)
        attributes,delimiter,mailbox = match.groups()
        # TODO: unquote mailbox (if necessary?)
        yield (attributes, delimiter, mailbox)
        children = True
        for attribute in ['\Noinferiors', '\HasNoChildren']:# RFC 2060 and 3348
            if attribute.lower() in attributes.lower():
                children = False
                break
        if not recursive or not children:
            continue
        for m in traverse_mailboxes(imap_connection, root=mailbox):
            yield m

def get_quota(imap_connection, mailbox='inbox'):
    status,value = imap_connection.getquotaroot(mailbox)
    #status,value = ('OK',
    #    [['inbox user/xyz12'], ['user/xyz12 (STORAGE 80000 102400)']])
    quota = value[1][0]
    regexp = re.compile(r'(.*) \(STORAGE ([0-9]*) ([0-9]*)\)')
    m = regexp.match(quota)
    root,used,avail = m.groups()
    return (root, int(used), int(avail))

def get_mailbox_size(imap_connection, mailbox='inbox'):
    status,value = imap_connection.select(mailbox, readonly=True)
    if status != 'OK':
        raise Exception(value)
    status,value = imap_connection.search(None, 'ALL')
    size = 0
    regexp = re.compile(r'.*RFC822\.SIZE ([0-9]*).*')
    for num in value[0].split():
        status,value = imap_connection.fetch(num, '(RFC822.SIZE)')
        #status,value = ('OK', ['231 (RFC822.SIZE 4882)'])
        match = regexp.match(value[0])
        msg_size = int(match.group(1))
        size += msg_size
    return size


if __name__ == '__main__':
    import optparse
    p = optparse.OptionParser('%prog SERVER[:port] account',
        epilog='imapquota -s imap.mail.drexel.edu xyz12@drexel.edu')
    p.add_option('-s', '--ssl', dest='ssl', default=False, action='store_true',
                 help='Connect over an SSL encrypted socket and change default port to 993.')
    p.add_option('-m', '--mailbox', dest='mailbox',
                 help='Root mailbox (e.g. INBOX).')
    p.add_option('-r', '--recursive', dest='recursive', default=False,
                 action='store_true',
                 help='List size of all mailboxes below the root given with `-m`.')
    options,args = p.parse_args()

    server,account = args
    server,port = parse_server(server, ssl=options.ssl)
    password = getpass.getpass()

    quota_roots = {}

    if options.ssl == True:
        i = imaplib.IMAP4_SSL(server, port)
    else:
        i = imaplib.IMAP4(server, port)
    i.login(account, password)
    try:
        if options.mailbox:
            if options.recursive:
                mailboxes = [
                    m for a,d,m in traverse_mailboxes(i, options.mailbox)]
                mailboxes.insert(0, options.mailbox)
            else:
                mailboxes = [options.mailbox]
        elif options.recursive:
            mailboxes = [m for a,d,m in traverse_mailboxes(i)]
        else:
            mailboxes = ['inbox']
        for mailbox in mailboxes:
            root,used,avail = get_quota(i, mailbox)
            if root not in quota_roots:
                quota_roots[root] = (used, avail)
            if options.recursive:
                size = get_mailbox_size(i, mailbox)
                print '%s used %d B' % (mailbox, size)
    finally:
        i.logout()

    for root,specs in sorted(quota_roots.iteritems()):
        used,available = specs
        print 'quota rooted at %s used %d of %d KB (%.1f percent)' % (
            root, used, avail, (float(used)/avail*100.0))