import-mailbox-to-gmail.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Import mbox files to a specified label for many users.

Liron Newman lironn@google.com

Copyright 2015 Google Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import argparse
import io
import logging
import logging.handlers
import mailbox
import os
import sys
import threading
from concurrent.futures import ThreadPoolExecutor

import httplib2
import oauth2client.tools
from apiclient import discovery
from googleapiclient.http import MediaIoBaseUpload, set_user_agent
from oauth2client.service_account import ServiceAccountCredentials
import OpenSSL  # Required by Google API library, but not checked by it

APPLICATION_NAME = 'import-mailbox-to-gmail'
APPLICATION_VERSION = '1.5'

SCOPES = ['https://www.googleapis.com/auth/gmail.insert',
          'https://www.googleapis.com/auth/gmail.labels']

parser = argparse.ArgumentParser(
    description='Import mbox files to a specified label for many users.',
    formatter_class=argparse.RawDescriptionHelpFormatter,
    parents=[oauth2client.tools.argparser],
    epilog=
    """
 * The directory needs to have a subdirectory for each user (with the full
   email address as the name), and in it there needs to be a separate .mbox
   file for each label. File names must end in .mbox.

 * Filename format: <user@domain.com>/<labelname>.mbox.
   Example: joe@mycompany.com/Migrated messages.mbox - This is a file named
   "Migrated messages.mbox" in the "joe@mycompany.com" subdirectory.
   It will be imported into joe@mycompany.com's Gmail account under the
   "Migrated messages" label.

 * See the README at https://goo.gl/JnFt0x for more usage information.
""")
parser.add_argument(
    '--json',
    required=True,
    help='Path to the JSON key file from https://console.developers.google.com')
parser.add_argument(
    '--dir',
    required=True,
    help=
    'Path to the directory that contains user directories with mbox files to '
    'import')
parser.add_argument(
    '--no-fix-msgid',
    dest='fix_msgid',
    required=False,
    action='store_false',
    help=
    "Don't fix Message-ID headers that are missing brackets "
    "(default: fix them)")
parser.add_argument(
    '--noreplaceqp',
    dest='replace_quoted_printable',
    required=False,
    action='store_false',
    help=
    "Replace 'Content-Type: text/quoted-printable' with text/plain (default: "
    "replace it)")
parser.add_argument(
    '--num_retries',
    default=10,
    type=int,
    help=
    'Maximum number of exponential backoff retries for failures (default: 10)')
parser.add_argument(
    '--log',
    required=False,
    default='%s-%d.log' % (APPLICATION_NAME, os.getpid()),
    help=
    'Optional: Path to a the log file (default: %s-####.log in the current '
    'directory, where #### is the process ID)' % APPLICATION_NAME)
parser.add_argument(
    '--httplib2debuglevel',
    default=0,
    type=int,
    help='Debug level of the HTTP library: 0=None (default), 4=Maximum.')
parser.add_argument(
    '--from_message',
    default=0,
    type=int,
    help=
    'Message number to resume from, affects ALL users and ALL '
    'mbox files (default: 0)')
parser.set_defaults(fix_msgid=True, replace_quoted_printable=True,
                    logging_level='INFO')
args = parser.parse_args()


def get_credentials(username):
    """Gets valid user credentials from a JSON service account key file.

    Args:
      username: The email address of the user to impersonate.
    Returns:
      Credentials, the obtained credential.
    """
    credentials = ServiceAccountCredentials.from_json_keyfile_name(
        args.json,
        scopes=SCOPES).create_delegated(username)

    return credentials


def get_label_id_from_name(service, username, labels, labelname):
    """Get label ID if it already exists, otherwise create it."""
    if labelname.endswith('.mbox'):
        # Strip .mbox suffix from folder names
        labelname = labelname[:-5]
    for label in labels:
        if label['name'].upper() == labelname.upper():
            return label['id']

    logging.info("Label '%s' doesn't exist, creating it", labelname)
    try:
        label_object = {
            'messageListVisibility': 'show',
            'name': labelname,
            'labelListVisibility': 'labelShow'
        }
        label = service.users().labels().create(
            userId=username,
            body=label_object).execute(num_retries=args.num_retries)
        logging.info("Label '%s' created", labelname)
        labels.append(label)
        return label['id']
    except Exception:
        logging.exception("Can't create label '%s' for user %s", labelname, username)
        raise


def process_mbox_files(username, service, labels, http_factory):
    """Iterates over the mbox files found in the user's subdir and imports them.

    Args:
      username: The email address of the user to import into.
      service: A Gmail API service object.
      labels: Dictionary of the user's labels and their IDs.
    Returns:
      A tuple of: Number of labels imported without error,
                  Number of labels imported with some errors,
                  Number of labels that failed completely,
                  Number of messages imported without error,
                  Number of messages that failed.
    """
    number_of_labels_imported_without_error = 0
    number_of_labels_imported_with_some_errors = 0
    number_of_labels_failed = 0
    number_of_messages_imported_without_error = 0
    number_of_messages_failed = 0
    base_path = os.path.join(args.dir, username)
    for root, dirs, files in os.walk(base_path):
        for dir in dirs:
            try:
                labelname = os.path.join(root[len(base_path) + 1:], dir)
                get_label_id_from_name(service, username, labels, labelname)
            except Exception:
                logging.error("Labels under '%s' may not nest correctly", dir)
        for file in files:
            filename = root[len(base_path) + 1:]
            if filename:
                filename += u'/'
            filename += file
            labelname, ext = os.path.splitext(filename)
            full_filename = os.path.join(root, file)
            if labelname.endswith('.mbox/mbox'):
                # Assume this is an Apple Mail export, so there's an mbox file inside a
                # dir that ends with .mbox.
                labelname = labelname[:-10]
                logging.info("File '%s' looks like an Apple Mail export, importing it "
                             "into label '%s'",
                             full_filename,
                             labelname)
            elif ext != '.mbox':
                logging.info("Skipping '%s' because it doesn't have a .mbox extension",
                             full_filename)
                continue
            if os.path.isdir(full_filename):
                # This "shouldn't happen" but it does, sometimes.
                # Assume this is an Apple Mail export, so there's an mbox file inside the dir.
                full_filename += os.path.join(full_filename, 'mbox')
                logging.info("Using '%s' instead of the directory", full_filename)
            logging.info("Starting processing of '%s'", full_filename)
            number_of_successes_in_label = 0
            number_of_failures_in_label = 0
            mbox = mailbox.mbox(full_filename)
            try:
                label_id = get_label_id_from_name(service, username, labels, labelname)
            except Exception:
                logging.error("Skipping label '%s' because it can't be created", labelname)
                continue
            logging.info("Using label name '%s', ID '%s'", labelname, label_id)

            def process_message(local, index, message):
                if index < args.from_message:
                    return
                logging.info("Processing message %d in label '%s'", index, labelname)
                try:
                    if (args.replace_quoted_printable and
                            'Content-Type' in message and
                            'text/quoted-printable' in message['Content-Type']):
                        message.replace_header(
                            'Content-Type', message['Content-Type'].replace(
                                'text/quoted-printable', 'text/plain'))
                        logging.info('Replaced text/quoted-printable with text/plain')
                except Exception:
                    logging.exception(
                        'Failed to replace text/quoted-printable with text/plain '
                        'in Content-Type header')
                try:
                    if args.fix_msgid and 'Message-ID' in message:
                        msgid = message['Message-ID']
                        if msgid[0] != '<':
                            msgid = '<' + msgid
                            logging.info('Added < to Message-ID: %s', msgid)
                        if msgid[-1] != '>':
                            msgid += '>'
                            logging.info('Added > to Message-ID: %s', msgid)
                        message.replace_header('Message-ID', msgid)
                except Exception:
                    logging.exception('Failed to fix brackets in Message-ID header')
                metadata_object = {'labelIds': [label_id]}
                try:
                    # Use media upload to allow messages more than 5mb.
                    # See https://developers.google.com/api-client-library/python/guide/media_upload
                    # and http://google-api-python-client.googlecode.com/hg/docs/epy/apiclient.http.MediaIoBaseUpload-class.html.
                    message_data = io.StringIO(message.as_string())
                    media = MediaIoBaseUpload(message_data, mimetype='message/rfc822')
                    message_response = service.users().messages().import_(
                        userId=username,
                        fields='id',
                        neverMarkSpam=True,
                        processForCalendar=False,
                        internalDateSource='dateHeader',
                        body=metadata_object,
                        media_body=media).execute(http=local.http, num_retries=args.num_retries)

                    logging.debug("Imported mbox message '%s' to Gmail ID %s",
                                  message.get_from(),
                                  message_response['id'])

                    return True
                except Exception:
                    logging.exception('Failed to import mbox message')
                    return False

            local = threading.local()

            # httplib2 is not thread safe
            def initializer_worker(local_context):
                local_context.http = http_factory()

            with ThreadPoolExecutor(initializer=initializer_worker, initargs=(local,), max_workers=10) as thread_pool:
                for result in thread_pool.map(lambda arguments: process_message(local, *arguments), enumerate(mbox)):
                    if result:
                        number_of_successes_in_label += 1
                    else:
                        number_of_failures_in_label += 1

            logging.info("Finished processing '%s'. %d messages imported "
                         "successfully, %d messages failed.",
                         full_filename,
                         number_of_successes_in_label,
                         number_of_failures_in_label)
            if number_of_failures_in_label == 0:
                number_of_labels_imported_without_error += 1
            elif number_of_successes_in_label > 0:
                number_of_labels_imported_with_some_errors += 1
            else:
                number_of_labels_failed += 1
            number_of_messages_imported_without_error += number_of_successes_in_label
            number_of_messages_failed += number_of_failures_in_label
    return (number_of_labels_imported_without_error,  # 0
            number_of_labels_imported_with_some_errors,  # 1
            number_of_labels_failed,  # 2
            number_of_messages_imported_without_error,  # 3
            number_of_messages_failed)  # 4


def main():
    """Import multiple users' mbox files to Gmail.

    """
    httplib2.debuglevel = args.httplib2debuglevel
    # Use args.logging_level if defined.
    try:
        logging_level = args.logging_level
    except AttributeError:
        logging_level = 'INFO'

    # Default logging to standard output
    logging.basicConfig(
        level=logging_level,
        format='%(asctime)s %(levelname)s %(funcName)s@%(filename)s %(message)s',
        datefmt='%H:%M:%S')

    # More detailed logging to file
    file_handler = logging.handlers.RotatingFileHandler(args.log,
                                                        maxBytes=1024 * 1024 * 32,
                                                        backupCount=8)
    file_formatter = logging.Formatter(
        '%(asctime)s %(process)d %(levelname)s %(funcName)s '
        '(%(filename)s:%(lineno)d) %(message)s')
    file_formatter.datefmt = '%Y-%m-%dT%H:%M:%S (%z)'
    file_handler.setFormatter(file_formatter)
    logging.getLogger().addHandler(file_handler)

    logging.info('*** Starting %s %s on Python %s ***',
                 APPLICATION_NAME,
                 APPLICATION_VERSION,
                 sys.version)
    logging.info('Arguments:')
    for arg, value in sorted(vars(args).items()):
        logging.info('\t%s: %r', arg, value)

    number_of_labels_imported_without_error = 0
    number_of_labels_imported_with_some_errors = 0
    number_of_labels_failed = 0
    number_of_messages_imported_without_error = 0
    number_of_messages_failed = 0
    number_of_users_imported_without_error = 0
    number_of_users_imported_with_some_errors = 0
    number_of_users_failed = 0

    def http_factory():
        return credentials.authorize(set_user_agent(
                    httplib2.Http(),
                    '%s-%s' % (APPLICATION_NAME, APPLICATION_VERSION)))

    for username in next(os.walk(args.dir))[1]:
        try:
            logging.info('Processing user %s', username)
            try:
                credentials = get_credentials(username)
                http = http_factory()
                service = discovery.build('gmail', 'v1', http=http)
            except Exception:
                logging.error("Can't get access token for user %s", username)
                raise

            try:
                results = service.users().labels().list(
                    userId=username,
                    fields='labels(id,name)').execute(num_retries=args.num_retries)
                labels = results.get('labels', [])
            except Exception:
                logging.error("Can't get labels for user %s", username)
                raise

            try:
                result = process_mbox_files(username, service, labels, http_factory)
            except Exception:
                logging.error("Can't process mbox files for user %s", username)
                raise
            if result[2] == 0 and result[4] == 0:
                number_of_users_imported_without_error += 1
            elif result[0] > 0 or result[3] > 0:
                number_of_users_imported_with_some_errors += 1
            else:
                number_of_users_failed += 1
            number_of_labels_imported_without_error += result[0]
            number_of_labels_imported_with_some_errors += result[1]
            number_of_labels_failed += result[2]
            number_of_messages_imported_without_error += result[3]
            number_of_messages_failed += result[4]
            logging.info('Done importing user %s. Labels: %d succeeded, %d with some '
                         'errors, %d failed. Messages: %d succeeded, %d failed.',
                         username,
                         result[0],
                         result[1],
                         result[2],
                         result[3],
                         result[4])
        except Exception:
            number_of_users_failed += 1
            logging.exception("Can't process user %s", username)
    logging.info("*** Done importing all users from directory '%s'", args.dir)
    logging.info('*** Import summary:')
    logging.info('    %d users imported with no failures',
                 number_of_users_imported_without_error)
    logging.info('    %d users imported with some failures',
                 number_of_users_imported_with_some_errors)
    logging.info('    %d users failed',
                 number_of_users_failed)
    logging.info('    %d labels (mbox files) imported with no failures',
                 number_of_labels_imported_without_error)
    logging.info('    %d labels (mbox files) imported with some failures',
                 number_of_labels_imported_with_some_errors)
    logging.info('    %d labels (mbox files) failed',
                 number_of_labels_failed)
    logging.info('    %d messages imported successfully',
                 number_of_messages_imported_without_error)
    logging.info('    %d messages failed\n',
                 number_of_messages_failed)
    if (number_of_messages_failed + number_of_labels_failed +
            number_of_users_failed > 0):
        logging.info('*** Check log file %s for detailed errors.', args.log)
    logging.info('Finished.\n\n')


if __name__ == '__main__':
    main()