# $Id: mime.py,v 1.14 2002/03/14 18:22:54 kjetilja Exp $

## System modules
import string

## Local modules
import pygmymultifile
import pygmymimetools


##
## Function find_all_parts ()
##
def find_all_parts(parts, new_parts=[]):
    for sm in parts:
        type = sm.gettype()
        if (type in ['multipart/alternative', 'multipart/mixed']):
            find_all_parts(sm.body, new_parts)
        else:
            new_parts.append(sm)            
    return new_parts


##
## Function decode_quoted_printable()
##
def decode_quoted_printable(txt):
    # The version in the python distribution is very slow and essentially
    # tries to mimic a regular expression parser in python which is a
    # bad idea when you can use regular expressions directly.
    import quote_pri
    return quote_pri.decode(txt)

def decode_base64(txt):
    from StringIO import StringIO
    output = StringIO()

    pygmymimetools.decode(StringIO(txt), output, 'base64')

    return output.getvalue()

##
##
## Message class, handling multipart messages
##
##
class Message(pygmymimetools.Message):

    # Constructor
    def __init__(self, fp, start, stop):
        pygmymimetools.Message.__init__(self, fp)
        self.start = start
        self.stop = stop

    # Return the message's header text as a string.  If an
    # argument is specified, it is used as a filter predicate to
    # decide which headers to return (its argument is the header
    # name converted to lower case).
    def getheadertext(self, pred = None):
        if not pred:
            return string.joinfields(self.headers, '')
        headers = []
        hit = 0
        for line in self.headers:
            if line[0] not in string.whitespace:
                i = string.find(line, ':')
                if i > 0:
                    hit = pred(string.lower(line[:i]))
            if hit: headers.append(line)
        return string.joinfields(headers, '')

    # Return the message's body text as string.  This undoes a
    # Content-Transfer-Encoding, but does not interpret other MIME
    # features (e.g. multipart messages).  To suppress to
    # decoding, pass a 0 as argument
    def getbodytext(self, decode = 1):
        self.fp.seek(self.startofbody)
        encoding = self.getencoding()
        if not decode or encoding in ('', '7bit', '8bit', 'binary') or \
           encoding not in ('quoted-printable', 'base64'):
            return self.fp.read()
        if encoding == 'quoted-printable':
            # Use our own handler here
            return decode_quoted_printable(self.fp.read())
        from cStringIO import StringIO
        output = StringIO()
        try:
            pygmymimetools.decode(self.fp, output, encoding)
        except:
            pass
        return output.getvalue()

    # Only for multipart messages: return the message's body as a
    # list of SubMessage objects.  Each submessage object behaves
    # (almost) as a Message object.
    def getbodyparts(self):
        if self.getmaintype() != 'multipart':
            raise SystemError, 'Content-Type is not multipart/*'
        bdry = self.getparam('boundary')
        if not bdry:
            raise SystemError, 'multipart/* without boundary param'
        self.fp.seek(self.startofbody)
        mf = pygmymultifile.MultiFile(self.fp, self.stop, 1)
        mf.push(bdry)
        parts = []
        while mf.next():
            parts.append( SubMessage(mf) )
        mf.pop()
        return parts

    # Return body, either a string or a list of messages
    def getbody(self):
        if self.getmaintype() == 'multipart':
            return self.getbodyparts()
        else:
            return self.getbodytext()


##
##
## SubMessage class, handle entries of a multipart message
##
##
class SubMessage(Message):
    # Constructor
    def __init__(self, fp):
        Message.__init__(self, fp, 0, 0)
        if self.getmaintype() == 'multipart':
            self.body = Message.getbodyparts(self)
        else:
            try:
                self.body = Message.getbodytext(self)
            except:
                self.body = self.bodyencoded = ''
                return
        self.bodyencoded = Message.getbodytext(self, decode=1)
        # If this is big, should remember file pointers

    def getbodytext(self, decode = 0):
        if not decode:
            return self.bodyencoded
        if type(self.body) == type(''):
            return self.body

    def getbodyparts(self):
        if type(self.body) == type([]):
            return self.body

    def getbody(self):
        return self.body