Add recursive mailbox size option to imapquota.py.

diff --git a/posts/imapquota/imapquota.py b/posts/imapquota/imapquota.py
index 92c0a74ef00956156448ec12605af6e48504b686..0f4a14b2384d44a959c979884ae531dc8363367e 100755 (executable)
--- a/posts/imapquota/imapquota.py
+++ b/posts/imapquota/imapquota.py
@@ -24,14 +24,54 @@ def parse_server(server, ssl=False):
         port = 143
     return (server, port)
 
-def get_quota(imap_connection):
-    status,value = imap_connection.getquotaroot('inbox')
-    #status,value = ('OK', [['inbox user/xyz12'], ['user/xyz12 (STORAGE 80000 102400)']])
+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]*).*')
+    regexp = re.compile(r'(.*) \(STORAGE ([0-9]*) ([0-9]*)\)')
     m = regexp.match(quota)
-    used,avail = [int(x) for x in m.groups()]
-    return (used, avail)
+    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__':
@@ -40,18 +80,47 @@ if __name__ == '__main__':
         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)
-    used,avail = get_quota(i)
-    i.logout()
-    print 'Used %d of %d KB' % (used, avail) 
-    print '%.1f percent' % (float(used)/avail*100.0)
+    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))