Connection error handling and some formatting

[imap-restore-mail] / file_iterator.py

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

Copyright (c) 2012 Intra2net AG
'''

import sys, os
import re

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

    def __init__(self):
        """Creates a connection and a user session."""
        self.created_mailboxes = []
        self.acl_mailboxes = []
        return

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

    def _message_read(self, filename):
        """Retrieves a message from the message file."""
        try:
            with open(filename, "r") as msgfile:
                message = msgfile.read()
        except:
            print("Could not open the e-mail file %s.", filename)
        return message

    def load_mailbox_list(self, mboxlistfile = ""):
        mboxdb = {}
        if mboxlistfile != "":
            try:
                with open(mboxlistfile, 'r') as acl_file:
                    for line in acl_file:
                        acls = {}
                        linedata = MBOXFILE_LINE.match(line).groups()
                        #!!! test this condition
                        if len(linedata) == 0:
                            print("Illegal line in mailbox list dump: %s" % line)
                            sys.exit()
                        key = linedata[0]
                        aclstr = linedata[1]
                        
                        # loop through acl rights string and build dictionary of users and rights
                        while(aclstr != ""):
                            acldata = ACL_STRING.match(aclstr).groups()
                            if len(acldata) == 0:
                                print("Illegal line in mailbox list dump: %s" % line)
                                sys.exit()                               
                            aclstr = acldata[2]
                            acls[acldata[0]]=acldata[1]

                        mboxdb[key] = acls

            except IOError:
                print("Could not open mboxlist file %s." % mboxlistfile)
        return mboxdb

    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)."""
        #print("Entered directory %s -> %s." % (filepath, mailpath))
        try:
            filepath = os.path.abspath(filepath)
            os.chdir(filepath)
        except OSError:
            print("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:
            #print("Now checking subpath %s in %s" % (subpath, filepath))
            if subpath == "." or subpath == "..":
                continue
            new_filepath = filepath + "/" + subpath
            if (os.path.isfile(new_filepath)):
                if os.path.getsize(new_filepath) == 0:
                    print("Skipping empty file %s." % subpath)
                else:
                    if MAIL_FILENAME.match(subpath):
                        #print("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:
                            print("Could not retrieve mail from file: %s" % new_filepath)                
            else:
                if os.path.isdir(new_filepath):
                    # cyrus ^ vs . storage replacement
                    subpath = subpath.replace("^", ".")
                    new_mailpath = mailpath + "/" + subpath
                    #print("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
                    #print("Done with directory %s and mailbox %s." % (new_filepath, new_mailpath))
        # mark mailboxes that need acl update
        self.acl_mailboxes.append(mailpath)
        return