path: root/distixlib/plugins/import_mail_plugin.py

# Copyright 2014  Lars Wirzenius
#
# 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.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# =*= License: GPL-3+ =*=


import contextlib
import email
import imaplib
import mailbox
import os
import re

import cliapp
import ttystatus

import distixlib


class WrongArguments(distixlib.StructuredError):

    msg = 'Wrong number of arguments: got {count}, wanted 2 or 3'


class NotKeywordArgument(distixlib.StructuredError):

    msg = '{arg} is not a keyword argumnent (KEY=VALUE)'


class ImportMailPlugin(cliapp.Plugin):

    commit_msg = 'imported mails'

    def enable(self):
        self.app.add_subcommand(
            'import-mail', self.import_mail,
            arg_synopsis='REPO FILE [KEY=VALUE]')

        self.app.add_subcommand(
            'import-mbox', self.import_mbox,
            arg_synopsis='REPO FILE [KEY=VALUE]')

        self.app.add_subcommand(
            'import-maildir', self.import_maildir,
            arg_synopsis='REPO MAILDIR')

        self.app.add_subcommand(
            'import-imap', self.import_imap,
            arg_synopsis='[KEY=VALUE]')

        self.app.settings.string(
            ['imap-server'],
            'which IMAP server to import from',
            metavar='ADDR')

        self.app.settings.string(
            ['imap-username'],
            'username on IMAP server')

        self.app.settings.string(
            ['imap-password-cmd'],
            'shell command to run to get password on IMAP server')

        self.app.settings.string(
            ['local-repos'],
            'where local repos should be when importing from IMAP')

    def import_mail(self, args):
        repo_dirname, mail_filename, keyvalue = self._parse_command_line(args)
        key, value = self._parse_keyvalue(keyvalue)
        msg = self._read_mail_message(mail_filename)

        context = _ImportContext()
        context.set_repo(distixlib.Repository(repo_dirname))
        context.repo.require_clean_working_tree()
        context.all_ticket_ids = context.store.get_ticket_ids()

        filenames = self._import_msg_to_ticket_store(context, msg, key, value)
        if filenames:
            context.repo.commit_changes(filenames, self.commit_msg)

    def import_mbox(self, args):
        self._import_folder(args, mailbox.mbox)

    def import_maildir(self, args):
        def maildir_factory(filename):
            return mailbox.Maildir(filename, factory=None)
        self._import_folder(args, maildir_factory)

    def import_imap(self, args):
        if args:
            key, value = self._parse_keyvalue(args[0])
        else:
            key, value = None, None

        collection = _RepoCollection(self.app.settings['local-repos'])

        imap_server = self.app.settings['imap-server']
        username = self.app.settings['imap-username']
        password_cmd = self.app.settings['imap-password-cmd']
        password = cliapp.runcmd(['sh', '-c', password_cmd]).strip()

        cp = self.app.settings.as_cp()
        repo_rules = _RepoRules()
        repo_rules.add_rules(cp)

        imap = imaplib.IMAP4_SSL(imap_server)
        imap.login(username, password)
        imap.select('INBOX')

        _, data = imap.search(None, 'ALL')
        for msgnum in data[0].split():
            print 'importing msg', msgnum
            _, data = imap.fetch(msgnum, '(RFC822)')
            _, text = data[0]
            repo_url = repo_rules.find_url(text)
            if repo_url:
                print '  to repo', repo_url
                msg = email.message_from_string(text)
                self._import_msg_into_repo(
                    collection, repo_url, msg, key, value)
                imap.store(msgnum, '+FLAGS', '(\\Deleted)')

        imap.expunge()
        imap.close()
        imap.logout()

    def _import_msg_into_repo(self, collection, repo_url, msg, key, value):
        if collection.is_local(repo_url):
            collection.pull(repo_url)
        else:
            collection.clone(repo_url)
        localdir = collection.localdir(repo_url)

        context = _ImportContext()
        context.set_repo(distixlib.Repository(localdir))
        context.repo.require_clean_working_tree()
        context.all_ticket_ids = context.store.get_ticket_ids()

        filenames = self._import_msg_to_ticket_store(context, msg, key, value)
        if filenames:
            context.repo.commit_changes(filenames, self.commit_msg)

        collection.push(repo_url)

    def _import_folder(self, args, folder_factory):
        repo_dirname, folder_filename, keyvalue = self._parse_command_line(
            args)
        key, value = self._parse_keyvalue(keyvalue)
        folder = folder_factory(folder_filename)

        context = _ImportContext()
        context.set_repo(distixlib.Repository(repo_dirname))
        context.repo.require_clean_working_tree()
        context.all_ticket_ids = context.store.get_ticket_ids()

        if self.app.settings['quiet']:
            progress = _QuietProgressReporter()
        else:
            progress = _MboxProgressReporter(len(folder))

        filenames = []
        with contextlib.closing(folder), progress:
            for msg in folder:
                progress.next_msg()
                filenames += self._import_msg_to_ticket_store(
                    context, msg, key, value)
        filenames += context.store.save_changes()
        if filenames:
            context.repo.commit_changes(filenames, self.commit_msg)

    def _parse_command_line(self, args):
        if len(args) == 2:
            return args[0], args[1], None
        elif len(args) == 3:
            if '=' not in args[2]:
                raise NotKeywordArgument(arg=args[2])
            return args[0], args[1], args[2]
        else:
            raise WrongArguments(count=len(args))

    def _parse_keyvalue(self, keyvalue):
        if keyvalue is None:
            return None, None
        return keyvalue.split('=', 1)

    def _read_mail_message(self, mail_filename):
        with open(mail_filename) as f:
            return email.message_from_file(f)

    def _import_msg_to_ticket_store(self, context, msg, key, value):

        referenced_ids = self._find_tickets_with_mails_referenced_by_msg(
            context, msg)
        msg_ids = distixlib.get_ids_from_message(msg)
        filenames = []

        if referenced_ids:
            for ticket_id in referenced_ids:
                if not self._contains_message(context.store, ticket_id, msg):
                    ticket = context.store.get_ticket(ticket_id)
                    ticket.add_message(msg)
                    self._set_key_value(ticket, key, value)
                    context.cache.add_msg_ids_for_ticket_id(
                        ticket.get_ticket_id(), msg_ids)
        else:
            if not self._is_already_imported(context, msg):
                new_ticket = self._create_ticket_from_msg(context.repo, msg)
                self._set_key_value(new_ticket, key, value)
                context.cache.add_msg_ids_for_ticket_id(
                    new_ticket.get_ticket_id(), msg_ids)
                context.all_ticket_ids.append(new_ticket.get_ticket_id())
                filenames = context.store.add_ticket(new_ticket)

        filenames += context.store.save_changes()
        return filenames

    def _set_key_value(self, ticket, key, value):
        if key is not None and value is not None:
            metadata = ticket.get_ticket_metadata()
            if key in metadata:
                metadata.remove_all_values(key)
            metadata.add(key, value)
            ticket.set_ticket_metadata(metadata)

    def _contains_message(self, store, ticket_id, msg):
        return store.ticket_has_message_with_text(
            ticket_id, msg.as_string())

    def _is_already_imported(self, context, msg):
        for ticket_id in context.all_ticket_ids:
            if self._contains_message(context.store, ticket_id, msg):
                return True
        return False

    def _find_tickets_with_mails_referenced_by_msg(self, context, msg):
        ticket_ids = []
        msg_ids = distixlib.get_ids_from_message(msg)

        for ticket_id in context.all_ticket_ids:
            other_ids = self._get_ticket_message_ids(context, ticket_id)
            if other_ids.intersection(msg_ids):
                ticket_ids.append(ticket_id)

        return ticket_ids

    def _get_ticket_message_ids(self, context, ticket_id):
        if ticket_id not in context.cache:
            filenames = context.store.get_message_filenames(ticket_id)
            for filename in filenames:
                msg = self._read_mail_message(filename)
                msg_ids = distixlib.get_ids_from_message(msg)
                context.cache.add_msg_ids_for_ticket_id(ticket_id, msg_ids)
        return context.cache.get_msg_ids_for_ticket_id(ticket_id)

    def _create_ticket_from_msg(self, repo, msg):
        ticket_id = repo.invent_new_ticket_id()
        subject = self._get_header(msg, 'Subject')
        ticket = self._create_ticket(ticket_id, subject)
        ticket.add_message(msg)
        return ticket

    def _get_header(self, msg, name):
        decoded = email.header.decode_header(msg[name])
        combined = u' '.join(
            self._safe_decode(value, encoding)
            for value, encoding in decoded)
        return combined

    def _safe_decode(self, text, encoding):
        try:
            return text.decode(encoding or 'us-ascii')
        except LookupError:
            return repr(text)
        except UnicodeDecodeError:
            return repr(text)

    def _create_ticket(self, ticket_id, title):
        ticket = distixlib.Ticket()
        ticket.set_ticket_id(ticket_id)
        ticket.set_title(title)
        return ticket


class _ImportContext(object):

    def __init__(self):
        self.repo = None
        self.store = None
        self.all_ticket_ids = None
        self.cache = _MessageIdCache()

    def set_repo(self, repo):
        self.repo = repo
        self.store = self.repo.open_ticket_store(distixlib.tickets_dir_name)


class _MessageIdCache(object):

    def __init__(self):
        self._dict = {}

    def __contains__(self, ticket_id):
        return ticket_id in self._dict

    def get_msg_ids_for_ticket_id(self, ticket_id):
        return self._dict.get(ticket_id, set())

    def add_msg_ids_for_ticket_id(self, ticket_id, msg_ids):
        old = self.get_msg_ids_for_ticket_id(ticket_id)
        self._dict[ticket_id] = old.union(msg_ids)


class _MboxProgressReporter(object):

    def __init__(self, total):
        self._ts = ttystatus.TerminalStatus()
        self._ts.format(
            '%ElapsedTime() '
            'importing message %Integer(current) of %Integer(total) '
            '(%PercentDone(current,total))')
        self._ts['current'] = 0
        self._ts['total'] = total

    def next_msg(self):
        self._ts['current'] += 1

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self._ts.clear()
        self._ts.finish()


class _QuietProgressReporter(object):

    def next_msg(self):
        pass

    def __enter__(self):
        pass

    def __exit__(self, *args):
        pass


class _RepoCollection(object):

    def __init__(self, locals_dir):
        self._locals_dir = locals_dir

    def is_local(self, repo_url):
        localdir = self.localdir(repo_url)
        return os.path.exists(localdir)

    def clone(self, repo_url):
        localdir = self.localdir(repo_url)
        cliapp.runcmd(['git', 'clone', repo_url, localdir])
        cliapp.runcmd(
            ['git', 'branch', '-u', 'origin', 'master'],
            cwd=localdir)

    def pull(self, repo_url):
        localdir = self.localdir(repo_url)
        cliapp.runcmd(['git', 'pull', '--rebase'], cwd=localdir)

    def push(self, repo_url):
        localdir = self.localdir(repo_url)
        print cliapp.runcmd(['git', 'push'], cwd=localdir)

    def localdir(self, url):
        dirname = self._dir_for_url(url)
        return os.path.join(self._locals_dir, dirname)

    def _dir_for_url(self, repo_url):
        s = repo_url
        s = '_'.join(s.split('/'))
        return s


class _RepoRules(object):

    def __init__(self):
        self._rules = []

    def add_rules(self, cp):
        for section in cp.sections():
            if section.startswith('distix:'):
                url = section[len('distix:'):].strip()
                pattern = cp.get(section, 'pattern')
                self._rules.append((url, pattern))

    def find_url(self, msg_text):
        msg_text = ''.join(msg_text.split('\r'))
        for url, pattern in self._rules:
            m = re.search(pattern, msg_text, re.M | re.I)
            if m is not None:
                return url
        return None