atomail.py

#!/usr/bin/env python
# coding=utf-8

################################################################################
# Module information
################################################################################

"""
  A package for converting mail into Atom (RFC 4287) RSS feeds.

  Mail messages can be retrieved from different sources, including
  stdin, NNTP, POP3, IMAP, and local mailboxes.

  For details on available parameters, use the --help option.
  Help and more information can be found on
    http://el-tramo.be/software/atomail
"""

__all__ = ['MessageFeed', 'MailSource', 'PipeSource', 'MailboxSource',
           'IMAPSource', 'POP3Source', 'NNTPSource', 'ATOM_NS']
__author__ = 'Remko Tronçon'
__version__ = '0.10'
__copyright__ = """
  Copyright (C) 2006-2016  Remko Tronçon

  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 2 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, write to the Free Software
  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
  """

# TODO:
# - Warn about missing SSL

################################################################################
# Imports
################################################################################

import sys
import os.path
import optparse
import datetime
import email
import email.header
import email.utils
import re
import xml
import string
import logging
import hashlib
import math
from xml.dom import minidom
import nntplib
import imaplib
import poplib
import mailbox
import html
import functools

################################################################################
# Constants
################################################################################

PROGRAM_NAME = 'AtoMail'
PROGRAM_URI = 'https://el-tramo.be/atomail'
PROGRAM_USAGESTRING = "usage: %prog [options] file"
PROGRAM_VERSIONSTRING = PROGRAM_NAME + ' ' + __version__ + '\nWritten by ' + \
    __author__ + '\n' + 'For more information, please visit ' + PROGRAM_URI

ATOM_NS = 'http://www.w3.org/2005/Atom'
# This will be used to decode headers if no encoding is specified. Should probably be smarter about this. See usage for more info
DEFAULT_ENCODING = "iso8859-1"

################################################################################
# Auxiliary functions and classes
################################################################################


class TZ(datetime.tzinfo):
    def __init__(self, seconds=0):
        self.seconds = seconds

    def utcoffset(self, dt):
        return datetime.timedelta(minutes=math.floor(self.seconds/60))


def message_id(message):
    hash = hashlib.md5()
    if message['From']:
        hash.update((message['From']).encode('utf-8'))
    if message['Subject']:
        hash.update((message['Subject']).encode('utf-8'))
    if message['Date']:
        hash.update((message['Date']).encode('utf-8'))
    return hash.hexdigest()


def message_date(message):
    date = email.utils.parsedate_tz(message["Date"])
    print(date)
    if date != None:
        return datetime.datetime(date[0], date[1], date[2], date[3], date[4], date[5], 0, TZ(date[9]))
    else:
        logging.warning('Unable to parse date \'' + message['Date'] + '\'')
        return datetime.datetime(datetime.MINYEAR, 1, 1)


def message_contents(message, default_charset):
    contents = []
    if message.is_multipart():
        for submessage in message.get_payload():
            contents += message_contents(submessage, default_charset)
    else:
        payload = message.get_payload(decode=True)
        if payload:
            content_type = message.get_content_type()
            if content_type == 'text/plain':
                contents = [('text', payload)]
            if content_type == 'text/html':
                contents = [('html', payload)]
        else:
            logging.warning('Missing payload in message')
    return contents


def current_datetime():
    now = datetime.datetime.now()
    current_tz = round(
        float((datetime.datetime.now() - datetime.datetime.utcnow()).seconds)/3600)
    if current_tz > 12:
        current_tz -= 24
    return now.replace(tzinfo=TZ(current_tz))


def from_atom_date(date):
    microseconds = 0
    timezone = 0
    suffix = date[19:]
    m = re.match("(\.(?P<microseconds>\d+))?(?P<timezone>(\+|-)\d\d)?", suffix)
    if m != None:
        if m.group('microseconds'):
            microseconds = int(m.group('microseconds'))
        if m.group('timezone'):
            timezone = int(m.group('timezone'))
    return datetime.datetime(int(date[0:4]), int(date[5:7]), int(date[8:10]), int(date[11:13]), int(date[14:16]), int(date[17:19]), microseconds, TZ(timezone))


def entry_date(entry):
    updated_nodes = entry.getElementsByTagName('updated')
    if updated_nodes:
        date = from_atom_date(updated_nodes[0].childNodes[0].data)
    else:
        date = datetime.datetime(datetime.MINYEAR, 1, 1)
    return date


def decode_header(header, default):
    decoded_header = ""
    for (result, encoding) in email.header.decode_header(header if header else default):
        # Some mails don't have an encoding in the header, yet they do encode the
        # header. We should do trial and error here, but for now assume it's
        # Iso8859-1.
        if not encoding:
            encoding = DEFAULT_ENCODING
        try:
            result = result.decode(
                encoding=encoding, errors='ignore') if encoding else result
        except AttributeError:
            pass

        decoded_header += result

    return decoded_header


def get_charset(message, default="ascii"):
    if message.get_content_charset():
        return message.get_content_charset()
    elif message.get_charset():
        return message.get_charset()
    else:
        return default

################################################################################
# MessageFeed
################################################################################


class MessageFeed:
    """A class representing an (Atom) feed of messages"""

    def __init__(self, filename, uri, title, max_items, max_time, strip_subject):
        """Initialize the feed"""
        logging.info('Initializing the message feed')
        self.filename = filename
        self.max_items = max_items
        self.max_time = max_time
        self.strip_subject = strip_subject
        self.doc = xml.dom.minidom.getDOMImplementation(
        ).createDocument(ATOM_NS, 'feed', None)
        self.doc.documentElement.setAttribute('xmlns', ATOM_NS)
        self.set_id(uri)
        self.set_updated(datetime.datetime(datetime.MINYEAR, 1, 1))
        if os.path.isfile(self.filename):
            logging.info('Reading feed from ' + self.filename)
            try:
                self.doc = minidom.parse(self.filename)
                self.doc.normalize()
            except:
                logging.warning('Unable to parse feed. Resetting the file.')
        else:
            logging.info('Creating new file ' + self.filename)
        self.set_link(uri)
        self.set_title(title)

    def __del__(self):
        """Destroy the feed"""
        self.doc.unlink()

    def add_message(self, message):
        """Adds a new message to the feed"""
        logging.info('Adding new message to the feed')

        # Create a new entry
        entry = self.doc.createElement('entry')

        # ID
        id = self.doc.createElement('id')
        id.appendChild(self.doc.createTextNode(
            self.id() + '#' + message_id(message)))
        entry.appendChild(id)

        # Author
        from_address = decode_header(message["From"], "Anonymous")
        (name, address) = email.utils.parseaddr(from_address)
        author = self.doc.createElement('author')
        author_name = self.doc.createElement('name')
        if name and address:
            author_name.appendChild(self.doc.createTextNode(name))
            author_email = self.doc.createElement('email')
            author_email.appendChild(self.doc.createTextNode(address))
            author.appendChild(author_email)
        else:
            author_name.appendChild(self.doc.createTextNode(from_address))
        author.appendChild(author_name)
        entry.appendChild(author)
        logging.debug('Author: ' + name + ' (' + address + ')')

        # Date
        date = message_date(message)
        published = self.doc.createElement("published")
        published.appendChild(self.doc.createTextNode(date.isoformat()))
        entry.appendChild(published)
        logging.debug('Published: ' + date.isoformat())

        # Updated
        updated = self.doc.createElement('updated')
        updated.appendChild(self.doc.createTextNode(
            current_datetime().isoformat()))
        entry.appendChild(updated)

        # Subject
        title = self.doc.createElement('title')
        title_text = decode_header(message["Subject"], "(No Subject)")
        if self.strip_subject:
            title_text = re.sub('\[[a-zA-Z0-9:_\. -]*\]\s*', '', title_text)
        title.appendChild(self.doc.createTextNode(title_text))
        entry.appendChild(title)
        logging.debug('Title: ' + title_text)

        # Content
        content = self.doc.createElement('content')
        contents = message_contents(message, get_charset(message))
        if contents:
            # Add the preferred content
            contents.sort()
            (content_type, content_text) = contents[0]
            content_txt = content_text.decode()

            # Replace text content with HTML counterpart (if there is none)
            # According to the spec, all whitespace in 'text' content can be collapsed. Since
            # mails are typically formatted, we don't want this, so we use preformatted
            if content_type == "text":
                content_type = "html"
                content_txt = "<pre>" + \
                    html.escape(content_txt) + "</pre>"

            content.setAttribute('type', content_type)
            content.appendChild(self.doc.createTextNode(content_txt))
        else:
            logging.warning('No valid contents found')
            content.appendChild(self.doc.createTextNode(title_text))
        entry.appendChild(content)

        # Add the entry to the feed
        self.doc.documentElement.appendChild(entry)

    def set_generator(self):
        """Updates the generator (program) of this feed"""
        generator = self.doc.createElement('generator')
        generator.appendChild(self.doc.createTextNode(PROGRAM_NAME))
        generator.setAttribute('version', __version__)
        generator.setAttribute('uri', PROGRAM_URI)
        self.replace_element(generator)

    def id(self):
        """Returns the unique identifier of this feed"""
        return list(filter(lambda x: x.parentNode == self.doc.documentElement, self.doc.documentElement.getElementsByTagName('id')))[0].childNodes[0].data

    def contains_message(self, message):
        id = message_id(message)
        feed_id = self.id()
        for entry in self.doc.getElementsByTagName('entry'):
            for entry_id_element in entry.getElementsByTagName('id'):
                entry_id = entry_id_element.childNodes[0].data
                if entry_id == feed_id + '#' + id:
                    return True
        return False

    def updated(self):
        """Returns the last time this feed was updated (as a datetime object)"""
        return from_atom_date(list(filter(lambda x: x.parentNode == self.doc.documentElement, self.doc.documentElement.getElementsByTagName('updated')))[0].childNodes[0].data)

    def set_id(self, id):
        """Sets the unique identifier of this feed"""
        self.set_element_text('id', id)

    def set_title(self, title):
        """Sets the title of this feed"""
        self.set_element_text('title', title)

    def set_link(self, uri):
        """Sets the URI of this feed"""
        link = self.doc.createElement('link')
        link.setAttribute('rel', 'self')
        link.setAttribute('href', uri)
        self.replace_element(link)

    def set_element_text(self, tagname, value):
        """Sets the text contents of a toplevel element of this feed to this string"""
        text = self.doc.createElement(tagname)
        text.appendChild(self.doc.createTextNode(value))
        self.replace_element(text)

    def replace_element(self, element):
        """Replaces a toplevel element in the feed"""
        # Try to find the relevant node node
        nodes = list(filter(lambda x: x.parentNode == self.doc.documentElement,
                            self.doc.documentElement.getElementsByTagName(element.tagName)))
        if nodes:
            self.doc.documentElement.replaceChild(element, nodes[0])

        # Create a new node (before the entries)
        entries = self.doc.documentElement.getElementsByTagName('entry')
        if len(entries) > 0:
            self.doc.documentElement.insertBefore(element, entries[0])
        else:
            self.doc.documentElement.appendChild(element)

    def trim_entries(self):
        """Removes all the redundant and outdated entries"""
        logging.info('Trimming entries')
        entries = self.doc.documentElement.getElementsByTagName('entry')
        # entries.sort(lambda x, y: cmp(entry_date(x), entry_date(y)))
        sorted(entries, key=entry_date)

        # Trim based on the maximum number of items
        if self.max_items > 0:
            while len(entries) > self.max_items:
                logging.debug('Removing redundant entry')
                self.doc.documentElement.removeChild(entries.pop(0))

        # Trim based on the maximum time elapsed
        if self.max_time > 0:
            max_datetime = current_datetime() - datetime.timedelta(minutes=self.max_time)
            while entries and entry_date(entries[0]) < max_datetime:
                logging.debug('Removing outdated entry')
                self.doc.documentElement.removeChild(entries.pop(0))

    def set_updated(self, time):
        """Sets the updated time as a datetime object."""
        self.set_element_text('updated', time.isoformat())

    def save(self):
        """Saves this feed to file.

        The generator is updated, and the elements are trimmed before the feed
        is saved"""
        logging.info('Saving feed')
        self.set_generator()
        self.trim_entries()
        logging.info('Writing feed to file ' + self.filename)
        out = open(self.filename, 'w')
        out.write(self.doc.toxml())
        out.close()


################################################################################
# MailSource
################################################################################

class MailSource:
    """An abstract class used to retrieve mail messages from a certain source.

    Subclasses must implement the messages() function, which retrieves the
    messages from the concrete source."""

    def messages(self):
        """Retrieve all messages (in reverse order)"""
        pass


class PipeSource(MailSource):
    """A class that retrieves messages from stdin"""

    def messages(self):
        logging.info('Reading message from stdin')
        return [email.message_from_file(sys.stdin)]


class MailboxSource(MailSource):
    """A class that retrieves messages from a mailbox file.

    The type of the mailbox file needs to be specified as a mailbox class.
    """

    def __init__(self, filename, type):
        logging.info('Initializing mailbox source (file=' +
                     filename + ',type=' + type.__name__ + ')')
        self.type = type
        self.filename = filename

    def messages(self):
        logging.info('Reading mails from ' + self.filename)
        # if self.type == mailbox.Maildir:
        #     file = self.filename
        # else:
        #     file = open(self.filename, 'r')
        # mbox = self.type(file, email.message_from_file)
        mbox = mailbox.mbox(self.filename)
        mails = []
        for (key, mail) in mbox.iteritems():
            mails.append(mail)
        # mail = mbox.next()
        # while mail != None:
        #     mails.append(mail)
        #     mail = mbox.next()
        mails.reverse()
        return mails


class POP3Source(MailSource):
    """A class that retrieves messages from a POP3 server."""

    def __init__(self, host, port, user, password, ssl=False):
        logging.info('Initializing POP3 client')
        if ssl:
            pop = poplib.POP3_SSL
        else:
            pop = poplib.POP3
        if port:
            self.pop = pop(host, port)
        else:
            self.pop = pop(host)
        logging.info('Authenticating')
        self.pop.user(user)
        self.pop.pass_(password)

    def messages(self):
        logging.info('Retrieving POP3 list')
        nb_messages = len(self.pop.list()[1])
        logging.debug(str(nb_messages) + ' messages waiting')
        while nb_messages > 0:
            message_text = ''
            for j in self.pop.retr(nb_messages)[1]:
                message_text += j + '\n'
            message = email.message_from_string(message_text)
            if message:
                yield message
            else:
                logging.warn('Unable to parse message:\n' + message_text)
            nb_messages -= 1


class NNTPSource(MailSource):
    """A class that retrieves messages from an NNTP server."""

    def __init__(self, host, port, group, user, password):
        logging.info('Initializing NNTP client')
        if port:
            self.nntp = nntplib.NNTP(host, port, user=user, password=password)
        else:
            self.nntp = nntplib.NNTP(host, user=user, password=password)
        self.group = group

    def messages(self):
        logging.info('Retrieving article list')
        result = self.nntp.group(self.group)
        first = int(result[2])
        last = int(result[3])
        while last >= first:
            logging.debug('Retrieving header of article ' + str(last))
            try:
                head = string.join(self.nntp.head(str(last))[3], '\n')
                message = email.message_from_string(head)
                body = string.join(self.nntp.body(str(last))[3], '\n')
                message = email.message_from_string(head + '\n' + body)
                if message:
                    yield message
                else:
                    logging.warn('Unable to parse message:\n' +
                                 head + '\n' + body)
            except nntplib.NNTPTemporaryError:
                pass
            last -= 1


class IMAPSource(MailSource):
    """A class that retrieves messages from an IMAP server."""

    def __init__(self, host, port, user, password, mailbox, ssl=False):
        logging.info('Initializing IMAP client')
        if ssl:
            imap = imaplib.IMAP4_SSL
        else:
            imap = imaplib.IMAP4
        if port:
            self.imap = imap(host, port)
        else:
            self.imap = imap(host)
        logging.info('Authenticating')
        self.imap.login(user, password)
        if mailbox:
            logging.info('Opening mailbox \'' + mailbox + '\'')
            self.imap.select(mailbox=mailbox)
        else:
            logging.info('Opening INBOX')
            self.imap.select()

    def messages(self):
        logging.info('Retrieving relevant message numbers')
        _, msgnums = self.imap.search(None, 'ALL')
        msg_numbers = msgnums[0].split()
        while msg_numbers:
            msg_number = msg_numbers.pop()
            logging.info('Fetching article ' + msg_number)
            typ, data = self.imap.fetch(msg_number, '(RFC822)')
            message = email.message_from_string(data[0][1])
            if message:
                yield message
            else:
                logging.warn('Unable to parse message:\n' + data[0][1])


################################################################################
# Main program
################################################################################

if __name__ == "__main__":
    # Parse the arguments
    parser = optparse.OptionParser(
        usage=PROGRAM_USAGESTRING, version=PROGRAM_VERSIONSTRING)
    parser.add_option('-q', '--quiet', help='Turn off logging', action='store_const',
                      dest='loglevel', const=logging.ERROR, default=logging.WARNING)
    parser.add_option('-v', '--verbose', help='Turn on verbose logging',
                      action='store_const', dest='loglevel', const=logging.INFO)
    parser.add_option('-d', '--debug', help='Turn on debug logging',
                      action='store_const', dest='loglevel', const=logging.DEBUG)
    parser.add_option('-l', '--logfile', metavar='FILE',
                      help='Send log to a file (instead of stderr)')
    parser.add_option('-m', '--mode', metavar='MODE', help='The mode in which this script should operate (pipe,mbox,maildir,pop3,pop3-ssl,imap,imap-ssl,nntp). Default: %default',
                      type='choice', choices=['pipe', 'mbox', 'maildir', 'pop3', 'pop3-ssl', 'imap', 'imap-ssl', 'nntp'], default='pipe')
    parser.add_option('-u', '--uri', metavar='URI',
                      help='The URI of the target feed')
    parser.add_option('-t', '--title', metavar='TITLE',
                      help='The title of the target feed', default='AtoMail feed')
    parser.add_option('', '--max-items', metavar='ITEMS',
                      help='The maximum number of items in the feed. Default: %default', type='int', default=10)
    parser.add_option('', '--max-time', metavar='MINUTES',
                      help='The maximum number of elapsed minutes for items in the feed', type='int', default=-1)
    parser.add_option('-s', '--strip-subject', action='store_true', dest='strip_subject',
                      default=False, help='Strip mailing-list headers from the subject')
    parser.add_option('-f', '--file', metavar='FILE',
                      help='The file or directory to read messages from (mbox,maildir)')
    parser.add_option('', '--host', metavar='HOST',
                      help='The host to receive messages from (pop3,pop3-ssl,imap,imap-ssl,nntp)')
    parser.add_option('', '--port', metavar='PORT', type='int',
                      help='The host port to receive messages from (pop3,pop3-ssl,imap,imap-ssl,nntp)')
    parser.add_option('', '--user', metavar='USERNAME',
                      help='The user to authenticate with (pop3,pop3-ssl,imap,imap-ssl,nntp)')
    parser.add_option('', '--password', metavar='PASSWORD',
                      help='The password to authenticate with (pop3,pop3-ssl,imap,imap-ssl,nntp)')
    #parser.add_option('','--mailbox', metavar='PASSWORD', help='The IMAP mailbox to open (imap)')
    parser.add_option('', '--group', metavar='GROUP',
                      help='The group from which to retrieve messages (nntp)')
    (options, args) = parser.parse_args()

    # Initialize the logger
    if options.logfile:
        handler = logging.FileHandler(options.logfile)
    else:
        handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter('%(levelname)s: %(message)s'))
    logging.getLogger().addHandler(handler)
    logging.getLogger().setLevel(options.loglevel)

    # Get the filename
    if len(args) != 1:
        sys.exit('Filename missing')
    filename = args[0]

    # Check hostname
    if options.mode in ('pop3', 'pop3-ssl', 'imap', 'imap-ssl', 'nntp') and not options.host:
        sys.exit('Host missing')

    if options.mode in ('pop3', 'pop3-ssl', 'imap', 'imap-ssl') and not options.user:
        sys.exit('Username missing')

    if options.mode in ('pop3', 'pop3-ssl', 'imap', 'imap-ssl') and not options.password:
        sys.exit('Password missing')

    if options.mode in ('mbox', 'maildir') and not options.file:
        sys.exit('Mailbox file or directory missing')

    if options.mode == 'nntp' and not options.group:
        sys.exit('Group missing')

    # Substitute defaults if necessary
    if not options.uri:
        logging.warning('Feed URI missing. Using default URI.')
        options.uri = 'http://example.com/' + os.path.basename(filename)

    # Initialize the feed
    feed = MessageFeed(filename=filename, uri=options.uri, title=options.title,
                       max_items=options.max_items, max_time=options.max_time, strip_subject=options.strip_subject)
    updated_time = feed.updated()
    current_time = current_datetime()
    logging.debug('Current time: ' + str(current_time))
    logging.debug('Feed last updated: ' + str(updated_time))

    # Initialize the message source
    logging.info('Initializing the message source')
    if options.mode == "mbox":
        source = MailboxSource(options.file, mailbox.mbox)
    elif options.mode == "maildir":
        source = MailboxSource(options.file, mailbox.Maildir)
    elif options.mode == "pop3":
        source = POP3Source(options.host, options.port,
                            options.user, options.password)
    elif options.mode == "pop3-ssl":
        source = POP3Source(options.host, options.port,
                            options.user, options.password, ssl=True)
    elif options.mode == "imap":
        source = IMAPSource(options.host, options.port,
                            options.user, options.password, None)
    elif options.mode == "imap-ssl":
        source = IMAPSource(options.host, options.port,
                            options.user, options.password, None, ssl=True)
    elif options.mode == "nntp":
        source = NNTPSource(options.host, options.port,
                            options.group, options.user, options.password)
    else:
        source = PipeSource()

    # Determine the date from which messages should be retrieved
    if options.max_time > 0:
        from_time = max(
            current_time - datetime.timedelta(minutes=options.max_time), updated_time)
    else:
        from_time = updated_time

    # Add available messages to the feed
    count = 0
    for message in source.messages():
        if feed.contains_message(message):
            logging.info('Message already in feed. Stopped retrieving.')
            break
        else:
            feed.add_message(message)
            count += 1
            if options.max_items > 0 and count > options.max_items:
                logging.info('Maximum # of items reached. Stopped retrieving.')
                break

    # Add the messages to the feed
    feed.set_updated(current_time)
    feed.save()