~vcs-imports/quotient/main

# -*- test-case-name: quotient.test.test_mimemessage -*-

# This module is part of the Quotient project and is Copyright 2003 Divmod:
# http://www.divmod.org/.  This is free software.  You can redistribute it
# and/or modify it under the terms of version 2.1 of the GNU Lesser General
# Public License as published by the Free Software Foundation.

import os
import quopri
import base64
import rfc822
import time
import itertools

from cStringIO import StringIO

from atop.filepile import symlink
from twisted.python.failure import Failure

from twisted.internet.error import ConnectionDone
from twisted.persisted.styles import Versioned

from atop.tpython import iterateInReactor
from atop.store import Item, Pool
from atop.powerup import Powerup, IPowerStation

from twisted.python import components

def unquote(st):
    if len(st) > 1:
        if st[0] == st[-1] == '"':
            return st[1:-1].replace('\\\\', '\\').replace('\\"', '"')
        if st.startswith('<') and st.endswith('>'):
            return st[1:-1]
    return st

class HeaderBodyParser:
    def __init__(self, part, parent):
        self.parent = parent
        self.parsingHeaders = 1
        self.prevheader = None
        self.prevvalue = None
        self.warnings = []
        self.part = part
        self.bodyMode = 'body'
        self.gotFirstHeader = False

    def close(self):
        if self.parent:
            self.parent.close()

    def startBody(self, linebegin, lineend):
        self.parsingHeaders = 0
        self.part.headersLength = linebegin - self.part.headersOffset
        self.part.bodyOffset = lineend

    def lineReceived(self, line, linebegin, lineend):
        if self.parsingHeaders:
            if not self.gotFirstHeader:
                self.part.headersOffset = linebegin
                self.gotFirstHeader = True
            return self.parseHeaders(line, linebegin, lineend)
        else:
            return self.parseBody(line, linebegin, lineend)

    def warn(self, text):
        self.warnings.append(text)

    def finishHeader(self):
        if self.prevheader is not None:
            self.part[self.prevheader] = self.prevvalue
        self.prevheader = self.prevvalue = None

    def parseHeaders(self, line, linebegin, lineend):
        if not line:
            self.finishHeader()
            self.startBody(linebegin, lineend)
            return self
        if line[0] in ' \t':
            self.prevvalue += '\n' + line
            return self
        h = line.split(': ', 1)
        if len(h) == 2:
            self.finishHeader()
            header, value = h
            self.prevheader = header
            self.prevvalue = value
        elif line and line[-1] == ':':
            # is this even a warning case?  need to read the rfc... -glyph
            self.prevheader = line[:-1]
            self.prevvalue = ''
        else:
            self.warn("perhaps a body line?: %r" % line)
            self.finishHeader()
            self.startBody(linebegin, lineend)
            self.lineReceived(line, linebegin, lineend)
        return self

    def parseBody(self, line, linebegin, lineend):
        return getattr(self, "parse_" + self.bodyMode)(line, linebegin, lineend)

class MIMEMessageParser(HeaderBodyParser):
    bodyFile = None
    def startBody(self, linebegin, lineend):
        HeaderBodyParser.startBody(self, linebegin, lineend)
        self.boundary = self._calcBoundary()
        if self.boundary:
            self.finalBoundary = self.boundary + '--'
            self.bodyMode = 'preamble'
            return
        ctyp = self.part['content-type']
        if ctyp and ctyp.split()[0].strip().lower() == 'message/rfc822':
            self.bodyMode = 'rfc822'
            return
        self.bodyMode = 'body'
        # self.bodyFile = self.part.getBody("wb")
        # ^ was only used for on-the-fly decoding

    def close(self):
        if self.bodyFile:
            self.bodyFile.close()
        HeaderBodyParser.close(self)

    def _calcBoundary(self):
        ctype = self.part['content-type']
        if ctype and ctype.strip().lower().startswith('multipart'):
            parts = ctype.split(';')
            for part in parts:
                ps = part.split('=', 1)
                if len(ps) == 2:
                    key, val = ps
                    key = key.strip().lower()
                    if key.lower() == 'boundary':
                        return '--' + unquote(val.strip())
            return None
        else:
            return None

    def parse_body(self, line, b, e):
        # TODO: on-the-fly decoding
        return self

    def parse_rfc822(self, line, b, e):
        np = self.subpart(parent=self, factory=MIMEMessageParser)
        np.lineReceived(line, b, e)
        return np

    def subpart(self, parent=None, factory=None):
        if parent is None:
            parent = self
        if factory is None:
            factory = MIMEPartParser
        newpart = self.part.newChild()
        nmp = factory(newpart, parent)
        return nmp

    def parse_preamble(self, line, b, e):
        if line.strip('\r\n') == self.boundary:
            self.bodyMode = 'nextpart'
            return self.subpart()
        return self

    def parse_nextpart(self, line, b, e):
        if line.strip('\r\n') == self.boundary:
            # If it's a boundary here, that means that we've seen TWO
            # boundaries, one right after another!  I can only assume that the
            # sub-human cretins who have thusly encoded their MIME parts are
            # attempting to convey the idea that the message *really* has a
            # part-break there...
            return self
        nmp = self.subpart()
        nmp.lineReceived(line, b, e)
        return nmp

    def parse_postamble(self, line, b, e):
        return self

class MIMEPartParser(MIMEMessageParser):
    def parseBody(self, line, linebegin, lineend):
        if line.strip('\r\n') == self.parent.boundary:
            # my body is over now - this is a boundary line so don't count it
            self.part.bodyLength = linebegin - self.part.bodyOffset
            return self.parent
        elif line == self.parent.finalBoundary:
            self.parent.bodyMode = 'postamble'
            self.part.bodyLength = linebegin - self.part.bodyOffset
            return self.parent
        else:
            return MIMEMessageParser.parseBody(self, line, linebegin, lineend)

    def parse_rfc822(self, line, linebegin, lineend):
        np = self.subpart(parent=self.parent)
        np.lineReceived(line, linebegin, lineend)
        return np

class MIMEPart:
    def __init__(self, parent=None):
        self.parent = parent
        self.children = []
        self.headers = []

    # for parser use only
    def setHeadersInfo(self, hoffset, hlength):
        self.headersInfo = hoffset, hlength

    def setBodyInfo(self, boffset, blength):
        self.bodyInfo = boffset, blength

    # email.Message compat: note non-coding-standard-compliant method names

    def walk(self):
        yield self
        for child in self.children:
            for part in child.walk():
                yield part

    def get_all(self, field, failObj):
        return self.get(field, failObj)

    def get_filename(self, failObj=None):
        return self.get_param('filename', failObj, 'content-disposition')

    def get_param(self, param, failObj=None, header='content-type', unquote=True):
        h = self[header]
        if not h:
            return failObj
        param = param.lower()
        for pair in [x.split('=', 1) for x in h.split(';')[1:]]:
            if pair[0].strip().lower() == param:
                r = len(pair) == 2 and pair[1].strip() or ''
                if unquote:
                    return mimeparser.unquote(r)
                return r
        return failObj

    def newChild(self):
        c = MIMEPart(self)
        self.children.append(c)
        return c

    # email.Message compat

    def __setitem__(self, key, val):
        self.headers.append((key, val))

    def __getitem__(self, key, failobj=None):
        for k,v in self.headers:
            if key.lower() == k.lower():
                return v
        return failobj
    get = __getitem__

    def __contains__(self, name):
        return not not self.get(name)

    def has_key(self, name):
        return name in self

    def items(self):
        return self.headers

    def get_charset(self):
        return None

    def get_type(self, failobj=None):
        return self.get('content-type', failobj)

    def get_payload(self, decode=False):
        """Get the message payload.
        """
        f = self.openFile()
        offt = self.bodyOffset
        leng = self.bodyLength
        f.seek(offt)
        data = f.read(leng)
        if decode:
            ctran = self['content-transfer-encoding']
            if ctran:
                ct = ctran.lower().strip()
                if ct == 'quoted-printable':
                    return quopri.decodestring(data)
                elif ct == 'base64':
                    return base64.decodestring(data)
                elif ct == '7bit':
                    return data
        return data

    def _uberparent(self):
        o = self
        while o.parent:
            o = o.parent
        return o

    def openFile(self):
        return open(self._uberparent().filename, 'rb')

    def get_default_type(self):
        return 'text/plain'

    def get_content_type(self):
        missing = object()
        value = self.get('content-type', missing)
        if value is missing:
            return self.get_default_type()  
        ctype = value.split(';', 1)[0].lower().strip()
        if ctype.count('/') != 1:
            return 'text/plain'  
        return ctype

    def get_content_maintype(self):
        ctype = self.get_content_type()
        return ctype.split('/')[0]

    def get_content_subtype(self):
        ctype = self.get_content_type()
        return ctype.split('/')[1]

    def get_main_type(self, failobj=None):
        """Return the message's main content type if present."""
        missing = object()
        ctype = self.get_type(missing)
        if ctype is missing:
            return failobj
        if ctype.count('/') != 1:
            return failobj
        return ctype.split('/')[0]

    def is_multipart(self):
        return bool(self.children)

    def getdate(self, name):
        data = self.get(name)
        if data:
            return rfc822.parsedate(data)

    def getHeaderParams(self, hdrname):
        ctype = self[hdrname]
        typeinfo = ctype.split(';')
        ctype = typeinfo[0].strip().lower()
        params = {}
        for t in typeinfo[1:]:
            kv = t.split('=', 1)
            if len(kv) == 2:
                k = kv[0].strip().lower()
                v = kv[1].strip().strip('"')
                params[k] = v
        return params

    def getAttachmentName(self):
        params = self.getHeaderParams("content-type")
        for fnk in 'name', 'filename':
            if params.has_key(fnk):
                return params[fnk]
        else:
            gtl = self.get_type().split(';')[0].lower()
            ext = {'text/html':  'html',
                   'text/plain': 'plain',
                   'image/jpeg': 'jpeg',
                   'image/png':  'png',
                   'image/gif':  'gif'}.get(gtl, 'bin')
            return 'Unknown.'+ext

    def inferType(self):
        """Infer a content-type.  This will attempt to do something with
        garbage data that isn't properly typed.
        """
        ctype = self['content-type']
        if not ctype:
            return 'text/plain'
        if ctype.lower().startswith("application/octet-stream"):
            self.getAttachmentName()
            ext = params['name'].strip().split(".")[-1]
            if exts.has_key(ext):
                return exts[ext]
        return ctype

    def getTypedParts(self, *types):
        for part in self.walk():
            # possible change: rather than get_content_type, use inferType to
            # catch parts which are malformed MIME-ly but still valid data.
            if part.get_content_type() in types:
                yield part

    def getAttachments(self):
        for part in self.walk():
            cd = part['content-disposition']
            if cd:
                cd = cd.split(';')[0].strip().lower()
                if cd == 'attachment':
                    yield part

    def keys(self):
        return [k for k, v in self.headers]

    # STUBBED METHODS: these will prevent spambayes et. al. from raising
    # exceptions, but we should look into how far we want to support them.
    
    def __delitem__(self, thing):
        pass

    def add_header(self, header, value):
        pass

    def get_charsets(self, failObj=None):
        return []

    # end stubbed methods


class MIMEMessage(MIMEPart, Item, Versioned):
    parent = None
    smtpInfo = None
    
    # A string describing how this message came to us
    receivedVia = None

    # Reference to the contact who sent this message
    contactRef = None

    # Don't call Item.__init__ - we don't want to initialize the item part of
    # ourselves until we're filled out enough to exist in the database.

    def assignIDs(self):
        mimeID = 0
        for part in self.walk():
            part.mimeID = mimeID
            mimeID += 1

    persistenceVersion = 1

    def upgradeToVersion1(self):
        self.assignIDs()

    def getPartByID(self, mimeID):
        w = self.walk()
        c = 0
        for p in w:
            if mimeID == c:
                return p
            c += 1

    def addToStore(self, store):
        self.assignIDs()
        Item.__init__(self, store)

    def index_name(self):
        if hasattr(self,'contact'):
            return self.contact.name
        else:
            return self['from']

    def index_subject(self):
        return self['subject']
    
    def index_date(self):
        return self.dateReceived

    def index_pop(self, pool):
        return self.storeID

    def getDisplayPart(self):
        return self.getTypedParts('text/plain','text/html','text/rtf').next()


# message started - headers begin (begin of line)

# headers ended - headers end (begin of line), body begins (end of line)

# boundary hit - body ends for previous child (begin of line) headers begin for
#      next child (end of line)

# "rfc822-begin" - headers begin for sub-rfc822-message

# subpart headers ended - headers end for child (begin of line), body begins
#      for child (end of line)

# subpart ended - body

# message ended (body ends)

class MIMEMessageReceiver:
    def __init__(self, avatar, deliver, trustDateHeaders=False):
        self.avatar = avatar
        self.deliver = deliver
        self.trustDateHeaders = trustDateHeaders
        self.done = False
        self.lineReceived = self.firstLineReceived

    def makeConnection(self, t):
        # rhg protocol
        self.bytecount = 0
        self.connectionMade()

    def connectionMade(self):
        self.message = MIMEMessage()
        self.file = self.avatar.newFile()
        #         self.message._currentsize = self.file.tell
        # ^ causes problems with pickle, obviously
        self.parser = MIMEMessageParser(self.message, None)

    def firstLineReceived(self, line):
        del self.lineReceived
        if line.startswith('From '):
            return
        return self.lineReceived(line)

    def lineReceived(self, line):
        linebegin = self.bytecount
        self.bytecount += (len(line) + 1)
        lineend = self.bytecount
        self.file.write(line+'\n')
        newParser = self.parser.lineReceived(line, linebegin, lineend)
        oldParser = self.parser
        if newParser is not oldParser:
            self.parser = newParser

    def connectionLost(self, reason):
        if self.done:
            return
        self.file.abort()

    def messageDone(self):
        self.done = True
        localNow = time.time()
        gmtDate = time.gmtime(localNow)
        self.parser.part.bodyLength = (self.bytecount - self.parser.part.bodyOffset)
        if self.trustDateHeaders:
            try:
                rdate = time.struct_time(rfc822.parsedate(self.message['received'].split(';')[-1]))
            except:
                rdate = gmtDate
        else:
            rdate = gmtDate
        self.message['x-divmod-processed'] = rfc822.formatdate(localNow)
        self.message.dateReceived = rdate

        def _():
            self.file.flush()
            size = self.file.tell()
            self.message.size = size
            self.message.addToStore(self.avatar)
            dplist = [str(x) for x in rdate[:3]]    # Y/M/D
            dplist.append(str(self.message.storeID))
            # store/avatarid/Y/M/D/msgid
            self.file.close(os.path.join(*dplist))
            self.message.filename = self.file.finalpath
            self.deliver(self.message)
        self.avatar.transact(_)

    # utility methods

    def feedFile(self, f):
        """Feed a file in.
        """
        return iterateInReactor(self._deliverer(f)).addCallback(
            lambda x: self.message)

    def feedString(self, s):
        """Feed a string in.
        """
        return self.feedFile(StringIO(s))

    def feedFileNow(self, f):
        for x in self._deliverer(f):
            pass
        return self.message

    def feedStringNow(self, s):
        return self.feedFileNow(StringIO(s))

    def _deliverer(self, f):
        self.makeConnection(None)
        try:
            while True:
                line = f.readline()
                if not line:
                    break
                line = line.strip('\r\n')
                self.lineReceived(line)
                yield None
        except:
            self.connectionLost(Failure())
            raise
        else:
            self.messageDone()
            self.connectionLost(Failure(ConnectionDone()))

class IMIMEDelivery(components.Interface):
    """I am a MIME delivery object.  I can wrap a storage avatar.
    """

    def createMIMEReceiver(self, trustReceivedHeaders):
        """Create a MIME receiver.  'trustReceivedHeaders' is an option to
        specify the primary date index: if it is True, it will use the last
        'Received' header.  If False, it will use the current time of the
        message's delivery.  This is dependent upon the delivery mechanism.
        For example, SMTP should NOT trustReceivedHeaders, because although the
        message was received from another mail server whose clock is probably
        correct, the message is not *finished* being 'received' until the
        system the user uses to check their mail has got its hands on it (in
        this case, us).  POP3, on the other hand, SHOULD trustReceivedHeaders,
        because the mail hosting system on the other end of the POP connection
        has arguably already added a header as to when it arrived at the
        address represented by the POP account.

        Finally, file imports should always, always trustReceivedHeaders,
        otherwise you will end up with a date index with all the imported
        messages clustered within 5 seconds of each other.
        """

class MIMEDeliverator:
    __implements__ = IMIMEDelivery
    def __init__(self, avatar, arrivalRef):
        self.avatar = avatar
        self.arrivalRef = arrivalRef

    def getArrivalRef(self):
        "get arrival reference"
        return self.arrivalRef

    def createMIMEReceiver(self, trustReceivedHeaders):
        arrivalPool = self.arrivalRef.getItem()
        return MIMEMessageReceiver(self.avatar, arrivalPool.addItem, trustReceivedHeaders)

class EmailPowerup(Powerup):
    def setUpPools(self, avatar):
        p = Pool(avatar, name='arrival')
        avatar.getRootPool().addItem(p)
        avatar.setComponent(IMIMEDelivery, MIMEDeliverator(avatar, p.referenceTo()))

exts = {
    "jpg": "image/jpeg",
    "jpeg": "image/jpeg",
    "png": "image/png",
    "gif": "image/gif"
    }