Improved formatting and separation in blocks
[imap-restore-mail] / file_iterator.py
1 '''
2 restore-mail-inject.py - Tool to inject mails via IMAP
3
4 Copyright (c) 2012 Intra2net AG
5 '''
6
7 import os
8 import re
9 import logging
10
11 MAIL_FILENAME = re.compile("^[0-9]+\.$")
12 MBOXFILE_LINE = re.compile("^(.*?)\t(?:\d )?default[\t ](.*)$")
13 ACL_STRING = re.compile("^(.*?)[\t ](.*?)\t(.*)$")
14
15 class FileIterator:
16     """This class iterates through the e-mail files."""
17
18     # class attributes
19     # mailboxes created during file traversal
20     created_mailboxes = None
21     # mailboxes to update during file traversal
22     acl_mailboxes = None
23
24     def __init__(self):
25         """Creates a connection and a user session."""
26
27         self.created_mailboxes = []
28         self.acl_mailboxes = []
29
30         return
31
32     @classmethod
33     def _message_read(cls, filename):
34         """Retrieves a message from the message file."""
35
36         try:
37             with open(filename, "r") as msgfile:
38                 message = msgfile.read()
39         except IOError:
40             logging.warning("Could not open the e-mail file %s", filename)
41             raise
42
43         return message
44
45     @classmethod
46     def load_mailbox_list(cls, mboxlistfile = ""):
47         """Load the list of mailboxes and acl rights for each from file."""
48
49         mboxdb = {}
50         if mboxlistfile != "":
51             try:
52                 with open(mboxlistfile, 'r') as acl_file:
53                     for line in acl_file:
54                         acls = {}
55                         linedata = MBOXFILE_LINE.match(line).groups()
56                         if len(linedata) == 0:
57                             logging.warning("Illegal line in mailbox list dump: %s", line)
58                             continue
59                         key = linedata[0]
60                         aclstr = linedata[1]
61                         
62                         # loop through acl rights string and build dictionary of users and rights
63                         while(aclstr != ""):
64                             acldata = ACL_STRING.match(aclstr).groups()
65                             if len(acldata) == 0:
66                                 logging.error("Illegal acl string in mailbox list dump: %s", line)
67                                 continue
68                             aclstr = acldata[2]
69                             acls[acldata[0]] = acldata[1]
70
71                         mboxdb[key] = acls
72             except IOError:
73                 logging.warning("Could not open mboxlist file %s", mboxlistfile)
74
75         return mboxdb
76
77     def load_mails(self, filepath, mailpath):
78         """Loads all e-mails from file hierarchy.
79         This recursive generator always returns a tuple of
80         the next found (e-mail, mailbox to store, internaldate)."""
81
82         logging.debug("Entered directory %s -> %s", filepath, mailpath)
83         try:
84             filepath = os.path.abspath(filepath)
85             os.chdir(filepath)
86         except OSError:
87             logging.warning("Can't open the directory %s", filepath)
88             return
89         # mark mailboxes that should be created
90         self.created_mailboxes.append(mailpath)
91         subpaths = os.listdir(filepath)
92         for subpath in subpaths:
93             if subpath == "." or subpath == "..":
94                 continue
95             new_filepath = filepath + "/" + subpath
96             if (os.path.isfile(new_filepath)):
97                 if os.path.getsize(new_filepath) == 0:
98                     logging.info("Skipping empty file %s", subpath)
99                 else:
100                     if MAIL_FILENAME.match(subpath):
101                         logging.info("Injecting file %s", subpath)
102                         try:
103                             message = self._message_read(new_filepath)
104                             # suggest file modification date for internaldate
105                             yield (message, mailpath, os.path.getmtime(new_filepath))
106                         except IOError:
107                             logging.warning("Could not retrieve mail from the file %s", new_filepath)
108             else:
109                 if os.path.isdir(new_filepath):
110                     # cyrus ^ vs . storage replacement
111                     subpath = subpath.replace("^", ".")
112                     new_mailpath = mailpath + "/" + subpath
113                     logging.debug("Inserting mails from directory %s into mailbox %s", new_filepath, new_mailpath)
114                     # load_mails($mboxdbref, $origuser, $targetuser)
115                     rcrs_generator = self.load_mails(new_filepath, new_mailpath)
116                     # you enter the generator in the for loop
117                     for rcr in rcrs_generator:
118                         yield rcr
119                     logging.debug("Done with directory %s and mailbox %s", new_filepath, new_mailpath)
120         # mark mailboxes that need acl update
121         self.acl_mailboxes.append(mailpath)
122
123         return