Added main function exit code which is the number of warnings during script execution
[imap-restore-mail] / restore_mail_inject.py
1 '''
2 restore-mail-inject.py - Tool to inject mails via IMAP
3
4 Copyright (c) 2012 Intra2net AG
5
6 This program relies on the following assumptions:
7 - IMAP hierarchy separator is / (unixhierarysep=yes)
8 - INBOX maps to user/[username]
9 - internal encoding of . etc. as of cyrus 2.2
10 '''
11 import logging
12 import argparse
13 from mail_iterator import MailIterator
14 from file_iterator import FileIterator
15 from warnings_handler import WarningsHandler
16
17 LOG_FILENAME = "restore_mail_inject.log"
18 LOG_FILE_LEVEL = logging.DEBUG
19 LOG_SHELL_LEVEL = logging.INFO
20 LOG_UNCLEAN_EXIT_LEVEL = logging.WARNING
21
22 def main():
23     """Main function."""
24
25     # prepare configuration
26     args = configure_args()
27     warnings_handler = prepare_logger()
28     logging.info("The module restore_mail_inject.py started with user %s, folder %s and source %s." %
29           (args.user, args.folder, args.srcdir))
30
31     # connect
32     session = MailIterator(args.user)
33     #session = MailIterator("/var/imap/socket/imap", "cyrus", "geheim")
34     storage = FileIterator()
35
36     # retrieve mailbox list from the mailbox list file
37     mailbox_list = storage.load_mailbox_list(args.mboxlistfile)
38
39     # delete olf IMAP folders if no append requested
40     if not args.append:
41         session.delete_mailboxes(args.folder)
42         if args.folder == "INBOX":
43             session.clear_inbox_acls(args.user)
44
45     # inject emails
46     path_generator = storage.load_mails(args.srcdir, args.folder)
47     for message, mailbox, date_modified in path_generator:
48
49         # mailboxes marked for creating and acl update
50         # add acls after all subfolders or their acls will be derived from parent folder
51         for new_mailbox in storage.created_mailboxes:
52             session.create_mailbox(new_mailbox)
53         for acl_mailbox in storage.acl_mailboxes:
54             session.add_acls(acl_mailbox, mailbox_list, args.ouser, args.user)
55         storage.created_mailboxes = []
56         storage.acl_mailboxes = []
57
58         session.inject_message(message, mailbox, date_modified)
59
60     # last iteration mailboxes in case root mailbox has no e-mails for injection
61     for new_mailbox in storage.created_mailboxes:
62         session.create_mailbox(new_mailbox)
63     for acl_mailbox in storage.acl_mailboxes:
64         session.add_acls(acl_mailbox, mailbox_list, args.ouser, args.user)
65
66     logging.info("Finished injecting mails. Exiting with code %s." % warnings_handler.detected_problems)
67     return warnings_handler.detected_problems
68
69 def configure_args():
70     """Configure arguments and return them."""
71     # parse arguments
72     parser = argparse.ArgumentParser(description="Tool to inject mails via IMAP.")
73     parser.add_argument('-u', '--username', dest='user', action='store',
74                         required=True, help='user to store mails to')
75     parser.add_argument('-f', '--foldername', dest='folder', action='store',
76                         default="INBOX", help='folder to store mails to - if not specified we overwrite INBOX')
77     parser.add_argument('-s', '--sourcedir', dest='srcdir', action='store',
78                         required=True, help='folder to read mail from')
79     parser.add_argument('-m', '--mboxlistfile', dest='mboxlistfile', action='store',
80                         default="", help='mboxlist file (flat file format) to read the ACLs from')
81     parser.add_argument('-o', '--ouser', dest='ouser', action='store',
82                         default="", help='name of the original user (=username if not specified)')
83     parser.add_argument('-a', '--append', dest='append', action='store_true',
84                         default=False, help="append mails, don't delete anything")
85     args = parser.parse_args()
86
87     if (args.folder != "INBOX"):
88         args.folder = "INBOX/" + args.folder
89     if (args.ouser == ""):
90         args.ouser = args.user
91
92     return args
93
94 def prepare_logger():
95     """Sets up the logging functionality"""
96
97     # reset the log
98     with open(LOG_FILENAME, 'w'):
99         pass
100
101     # add basic configuration
102     logging.basicConfig(filename=LOG_FILENAME,
103                         format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
104                         level=LOG_FILE_LEVEL)
105
106     # add a handler for a console output
107     default_logger = logging.getLogger('')
108     console = logging.StreamHandler()
109     console.setLevel(LOG_SHELL_LEVEL)
110     formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
111     console.setFormatter(formatter)
112     default_logger.addHandler(console)
113     
114     # add a handler for warnings counting
115     warnings_handler = WarningsHandler()
116     warnings_handler.setLevel(LOG_UNCLEAN_EXIT_LEVEL)
117     default_logger.addHandler(warnings_handler)
118     return warnings_handler
119
120 if __name__ == "__main__":
121     main()