Regex for the optional folder argument and skip shared folders option by default

[imap-mark-seen] / src / mail_iterator.py

'''
The module contains the MailIterator class.

Copyright (c) 2012 Intra2net AG
Author: Plamen Dimitrov

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.
'''

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 IMAP 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
    # skip shared folders
    skip_shared_folders = None

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

        self.skip_shared_folders = skip_shared_folders

        # connect to server
        try:
            self.mail_con = imaplib.IMAP4_SSL(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])
            # detect if mailbox is shared and if skip flag is set iterate further
            if(self.skip_shared_folders and mailbox[2].split(mailbox[1])[0] == '"user'):
                logging.info("Mailbox %s is shared and therefore skipped.", mailbox[2])
                continue
            # 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.debug("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()