Add mmassive e-mail update for performance by using message set STORE command

[imap-mark-seen] / mail_iterator.py

'''
imap-mark-seen.py - Tool to mark all e-mails as seen

Copyright (c) 2012 Intra2net AG
Author: Plamen Dimitrov and Thomas Jarosch
'''

import sys
import imaplib, socket
import re
import logging

MAILBOX_RESP = re.compile(r'\((?P<flags>.*?)\) "(?P<delimiter>.*)" (?P<name>.*)')

#imaplib.Debug = 4

class MailIterator:
    """This class communicates with the e-mail server."""

    # class attributes
    # IMAP4_SSL for connection with an IMAP server
    mail_con = None
    # list of tuples (uidvalidity, mailboxname) for the retrieved mailboxes
    mailboxes = None
    # logged in status
    logged_in = None

    def __init__(self, server, username, password):
        """Creates a connection and a user session."""

        # connect to server
        try:
            self.mail_con = imaplib.IMAP4(server)
            logging.info("Connected to %s", server)
        except socket.error as ex:
            logging.error("Could not connect to host: %s", ex)
            sys.exit()

        # log in
        try:
            self.mail_con.login(username, password)
            self.logged_in = True
            logging.info("Logged in as %s.", username)
        except self.mail_con.error as ex:
            self.logged_in = False
            logging.error("Could not log in as user %s: %s", username, ex)
            sys.exit()

        # list mailboxes
        try:
            _result, mailboxes = self.mail_con.list()
        except self.mail_con.error as ex:
            logging.warning("Could not retrieve mailboxes for user %s: %s", username, ex)
        self.mailboxes = []
        for mailbox in mailboxes:
            mailbox = MAILBOX_RESP.match(mailbox.decode('iso-8859-1')).groups()
            self.mailboxes.append(mailbox)
        self.mailboxes = sorted(self.mailboxes, key=lambda box: box[2], reverse=True)

        return

    def __del__(self):
        """Closes the connection and the user session."""

        if(self.logged_in):
            self.mail_con.close()
            self.mail_con.logout()

    def __iter__(self):
        """Iterates through all mailboxes, returns (children,delimiter,name)."""

        for mailbox in self.mailboxes:
            logging.debug("Checking mailbox %s", mailbox[2])
            # select mailbox if writable
            try:
                self.mail_con.select(mailbox[2])
                logging.info("Processing mailbox %s", mailbox[2])
            except self.mail_con.readonly:
                logging.warning("Mailbox %s is not writable and therefore skipped", mailbox[2])
                continue
            yield mailbox

    def fetch_messages(self):
        """Fetches the messages from the current mailbox, return list of uids."""

        try:
            # Work around unsolicited server responses in imaplib by clearing them
            self.mail_con.response('SEARCH')
            _result, data = self.mail_con.uid('search', None, "ALL")
        except (self.mail_con.error):
            raise UserWarning("Could not fetch messages")
        mailid_list = data[0].split()
        return mailid_list

    def set_seen_messages(self, mid_range):
        """Sets the \\Seen flag for all messages with the respective mids."""
        
        try:
            # Work around unsolicited server responses in imaplib by clearing them
            self.mail_con.response('STORE')
            _result, data = self.mail_con.uid('STORE', mid_range, '+FLAGS', "(\Seen)")
            logging.info("New flags for messages %s are %s", mid_range, data)
        except (self.mail_con.error) as ex:
            raise UserWarning("Could not set the flags for some messages: %s", ex)
        self.mail_con.expunge()