Additional fix from the previous refactoring

[imap-restore-mail] / src / file_iterator.py

'''
restore-mail-inject.py - Tool to inject mails via IMAP

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

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 os
import re
import logging

MAIL_FILENAME = re.compile("^[0-9]+\.$")
MBOXFILE_LINE = re.compile("^(.*?)\t(?:\d )?default[\t ](.*)$")
ACL_STRING = re.compile("^(.*?)[\t ](.*?)\t(.*)$")

class FileIterator:
    """This class iterates through the e-mail files."""

    # class attributes
    # mailboxes created during file traversal
    created_mailboxes = None
    # mailboxes to update during file traversal
    acl_mailboxes = None
    # acls retrieved from a file
    file_acls = None

    def __init__(self):
        """Creates a connection and a user session."""

        self.created_mailboxes = []
        self.acl_mailboxes = []
        self.file_acls = {}

    @classmethod
    def _message_read(cls, filename):
        """Retrieves a message from the message file."""

        try:
            with open(filename, "r") as msgfile:
                message = msgfile.read()
        except IOError:
            logging.warning("Could not open the e-mail file %s", filename)
            raise

        return message

    def load_mailbox_list(self, mboxlistfile, original_user):
        """Load the list of mailboxes and acl rights for each from file."""

        try:
            with open(mboxlistfile, 'r') as acl_file:
                for line in acl_file:

                    # read a line using regex
                    acls = {}
                    try:
                        linedata = MBOXFILE_LINE.match(line).groups()
                    except AttributeError:
                        logging.warning("Illegal line in mailbox list dump: %s", line)
                        continue
                    aclstr = linedata[1]

                    # changes acls key encoding to internal cyrus format and makes it absolute (as folder).
                    key = linedata[0].replace("INBOX/", "user/" + original_user + "/")
                    key = key.replace(".", "^")
                    key = key.replace("/", ".")

                    # loop through acl rights string and build dictionary of users and rights
                    while(aclstr != ""):
                        try:
                            acldata = ACL_STRING.match(aclstr).groups()
                        except AttributeError:
                            logging.warning("Illegal acl string in mailbox list dump: %s", line)
                            aclstr = ""
                            continue
                        aclstr = acldata[2]
                        acls[acldata[0]] = acldata[1]

                    self.file_acls[key] = acls
        except IOError:
            logging.warning("Could not open mboxlist file %s", mboxlistfile)

    def load_mails(self, filepath, mailpath):
        """Loads all e-mails from file hierarchy.
        This recursive generator always returns a tuple of
        the next found (e-mail, mailbox to store, internaldate)."""

        logging.debug("Entered directory %s -> %s", filepath, mailpath)
        try:
            filepath = os.path.abspath(filepath)
            os.chdir(filepath)
        except OSError:
            logging.warning("Can't open the directory %s", filepath)
            return
        # mark mailboxes that should be created
        self.created_mailboxes.append(mailpath)
        
        subpaths = os.listdir(filepath)
        for subpath in subpaths:
            if subpath == "." or subpath == "..":
                continue
            new_filepath = filepath + "/" + subpath

            # if path is file validate name and inject
            if (os.path.isfile(new_filepath)):
                if os.path.getsize(new_filepath) == 0:
                    logging.info("Skipping empty file %s", subpath)
                else:
                    if MAIL_FILENAME.match(subpath):
                        logging.info("Injecting file %s", subpath)
                        try:
                            message = self._message_read(new_filepath)
                            # suggest file modification date for internaldate
                            yield (message, mailpath, os.path.getmtime(new_filepath))
                        except IOError:
                            logging.warning("Could not retrieve mail from the file %s", new_filepath)
            else:

                # if path is directory do recursive call
                if os.path.isdir(new_filepath):
                    # cyrus ^ vs . storage replacement
                    subpath = subpath.replace("^", ".")
                    new_mailpath = mailpath + "/" + subpath
                    logging.debug("Inserting mails from directory %s into mailbox %s", new_filepath, new_mailpath)
                    # load_mails($mboxdbref, $origuser, $targetuser)
                    rcrs_generator = self.load_mails(new_filepath, new_mailpath)
                    # you enter the generator in the for loop
                    for rcr in rcrs_generator:
                        yield rcr
                    logging.debug("Done with directory %s and mailbox %s", new_filepath, new_mailpath)

        # mark mailboxes that need acl update
        self.acl_mailboxes.append(mailpath)