Server added to command line arguments

[imap-restore-mail] / src / imap_restore_mail.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 logging
import argparse, getpass
from mail_iterator import MailIterator
from file_iterator import FileIterator
from warnings_handler import WarningsHandler

# logging settings
LOG_FILENAME = "imap_restore_mail.log"
LOG_FILE_LEVEL = logging.DEBUG
LOG_SHELL_LEVEL = logging.INFO
LOG_UNCLEAN_EXIT_LEVEL = logging.WARNING

def main():
    """Main function."""

    # prepare configuration
    args = configure_args()
    warnings_handler = prepare_logger()
    logging.info("The module restore_mail_inject.py started with user %s, destination %s and source %s.",
          args.user, args.destination, args.source)

    # connect to unix socket or server
    psw = getpass.getpass()
    session = MailIterator(args.server, args.user, psw)
    storage = FileIterator()

    # retrieve mailbox list from the mailbox list file
    mailbox_list = storage.load_mailbox_list(args.mboxlistfile)

    # delete olf IMAP folders if no append requested
    if not args.append:
        session.delete_mailboxes(args.destination)
        if args.destination == "INBOX":
            session.clear_inbox_acls(args.user)

    # inject emails
    path_generator = storage.load_mails(args.source, args.destination)
    for message, mailbox, date_modified in path_generator:

        # mailboxes marked for creating and acl update
        # add acls after all subfolders or their acls will be derived from parent folder
        for new_mailbox in storage.created_mailboxes:
            session.create_mailbox(new_mailbox)
        for acl_mailbox in storage.acl_mailboxes:
            session.add_acls(acl_mailbox, mailbox_list, args.ouser, args.user)
        storage.created_mailboxes = []
        storage.acl_mailboxes = []

        session.inject_message(message, mailbox, date_modified)

    # last iteration mailboxes in case root mailbox has no e-mails for injection
    for new_mailbox in storage.created_mailboxes:
        session.create_mailbox(new_mailbox)
    for acl_mailbox in storage.acl_mailboxes:
        session.add_acls(acl_mailbox, mailbox_list, args.ouser, args.user)

    logging.info("Finished injecting mails. Exiting with code %s.", warnings_handler.detected_problems)
    return warnings_handler.detected_problems

def configure_args():
    """Configure arguments and return them."""

    # parse arguments
    parser = argparse.ArgumentParser(description="Tool to inject mails via IMAP.")
    parser.add_argument('-d', '--directory', dest='source', action='store',
                        required=True, help='file directory to read mails from')
    parser.add_argument('-s', '--server', dest='server', action='store',
                        default="localhost", help='imap server name with default localhost')    
    parser.add_argument('-u', '--username', dest='user', action='store',
                        required=True, help='user to store mails to')
    parser.add_argument('-o', '--ouser', dest='ouser', action='store',
                        default="", help='name of the original user (=username if not specified)')
    parser.add_argument('-f', '--foldername', dest='destination', action='store',
                        default="INBOX", help='folder to store mails to - if not specified we overwrite INBOX')
    parser.add_argument('-m', '--mboxlistfile', dest='mboxlistfile', action='store',
                        default="", help='mboxlist file (flat file format) to read the ACLs from')
    parser.add_argument('-a', '--append', dest='append', action='store_true',
                        default=False, help="append mails, don't delete anything")
    args = parser.parse_args()

    # some default post processing of arguments
    if (args.destination != "INBOX"):
        args.destination = "INBOX/" + args.destination
    if (args.ouser == ""):
        args.ouser = args.user

    return args

def prepare_logger():
    """Sets up the logging functionality"""

    # reset the log
    with open(LOG_FILENAME, 'w'):
        pass

    # add basic configuration
    logging.basicConfig(filename=LOG_FILENAME,
                        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                        level=LOG_FILE_LEVEL)

    # add a handler for a console output
    default_logger = logging.getLogger('')
    console = logging.StreamHandler()
    console.setLevel(LOG_SHELL_LEVEL)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    console.setFormatter(formatter)
    default_logger.addHandler(console)

    # add a handler for warnings counting
    warnings_handler = WarningsHandler()
    warnings_handler.setLevel(LOG_UNCLEAN_EXIT_LEVEL)
    default_logger.addHandler(warnings_handler)

    return warnings_handler

if __name__ == "__main__":
    main()