########################################
# interface to mail server transfers
# designed to be mixed-in to subclasses
# which use the methods defined here;
# class allows askPopPassword to differ
# loads all if server doesn't do top
# doesn't handle threads or UI here
# progress callback funcs get status
# all calls raise exceptions on error
########################################
#???? saveparts: open in text mode for text/ contypes?

import mailconfig
import poplib, smtplib, os, mimetypes
import email.Parser, email.Utils, email.Encoders

from email.Message import Message
from email.MIMEMultipart import MIMEMultipart
from email.MIMEAudio import MIMEAudio
from email.MIMEImage import MIMEImage
from email.MIMEText import MIMEText
from email.MIMEBase import MIMEBase

class MailTool:                    # super class for all mail tools
    def trace(self, message):      # redef me to disable or log to file
        print message


####################################
# parsing and attachment extraction
####################################


class MailParser(MailTool):
    """
    methods for parsing message text, attachements

    subtle thing: Message object payloads are either a simple
    string for non-multipart messages, or a list of Message
    objects if multipart (possiby nested); we don't need to
    distinguish between the two cases here, because the Message
    walk generator always returns self first, and so works fine
    on non-multipart messages too (a single object is walked);

    for simple messages, the message body is always considered
    here to be the sole part of the mail;  for miltipart messages,
    the parts list includes the main message text, as well as all
    attachments;  this allows simple messages not of type text to
    be handled like attachments in a UI (e.g., saved, opened).
    """
    
    def saveParts(self, savedir, message):
        """
        store all parts of a message as files in a local directory;
        returns [('maintype/subtype', 'filename')] list for use by
        callers, but does not open any parts or attachments here;
        """
        if not os.path.exists(savedir):
            os.mkdir(savedir)
        partfiles = []
        for (ix, part) in enumerate(message.walk()):    # walk includes message
            maintype = part.get_content_maintype()      # ix includes multiparts 
            if maintype == 'multipart':
                continue                                # multipart/*: container
            else:                                       # else save part to file
                filename, contype = self.partName(part, ix)
                fullname = os.path.join(savedir, filename)
                fileobj  = open(fullname, 'wb')             # use binary mode
                fileobj.write(part.get_payload(decode=1))   # decode base64,etc
                fileobj.close()
                partfiles.append((contype, fullname))       # for caller to open
        return partfiles
    
    def partsList(self, message):
        """"
        return a list of filenames for all parts of an
        already-parsed message, using same file name logic
        as saveParts, but do not store the part files here
        """
        partfiles = []
        for (ix, part) in enumerate(message.walk()):
            maintype = part.get_content_maintype()
            if maintype == 'multipart':
                continue
            filename, contype = self.partName(part, ix)
            partfiles.append(filename)
        return partfiles

    def partName(self, part, ix):
        """
        extract filename and content type from message part;
        filename: tries Content-Disposition, then Content-Type
        name param, or generates one based on mimetype guess;
        """
        filename = part.get_filename()                # filename in msg hdrs?
        contype  = part.get_content_type()            # lower maintype/subtype
        if not filename:
            filename = part.get_param('name')         # try content-type name
        if not filename:
            if contype == 'text/plain':               # hardcode plain text ext
                ext = '.txt'                          # else guesses .ksh!
            else:
                ext = mimetypes.guess_extension(contype)
                if not ext: ext = '.bin'              # use a generic default
            filename = 'part-%03d%s' % (ix, ext)
        return (filename, contype)

    def findMainText(self, message):
        """
        for text-oriented clients, return the first text part;
        for the payload of a simple message, or all parts of
        a multipart message, looks for text/plain, then text/html,
        then text/*, before deducing that there is no text to
        display;  this is a heuristic, but covers most simple,
        multipart/alternative, and multipart/mixed messages;
        content-type defaults to text/plain if none in simple msg;        

        handles message nesting at top level by walking instead
        of list scans;  if non-multipart but type is text/html,
        returns the htlm as the text with an html type: caller
        may open in webbrowser;  if non-multipart and not text,
        no text to display: save/open in UI;  caveat: does not
        try to concatenate multiple inline text/plain parts
        """
        # try to find a plain text
        for part in message.walk():                        # walk visits messge
            type = part.get_content_type()                 # if non-multipart
            if type == 'text/plain':
                return type, part.get_payload(decode=1)    # may be base64

        # try to find a html part
        for part in message.walk():
            type = part.get_content_type()
            if type == 'text/html':
                return type, part.get_payload(decode=1)    # caller renders

        # try any other text type, including xml
        for part in message.walk():
            if part.get_content_maintype() == 'text':
                return part.get_content_type(), part.get_payload(decode=1)

        # punt: could use first part, but it's not marked as text 
        return 'text/plain', '[No text to display]'

    # returned when parses fail
    errorMessage = Message()
    errorMessage.set_payload('[Unable to parse message - format error]')
    
    def parseHeaders(self, mailtext):
        """
        parse headers only, return root email.Message object
        stops after headers parsed, even if nothing else follows (top)
        email.Message object is a mapping for mail header fields
        payload of message object is None, not raw body text
        """
        try:
            return email.Parser.Parser().parsestr(mailtext, headersonly=True)
        except:
            return self.errorMessage

    def parseMessage(self, fulltext):
        """
        parse entire message, return root email.Message object
        payload of message object is a string if not is_multipart()
        payload of message object is more Messages if multiple parts
        the call here same as calling email.message_from_string()
        """
        try:
            return email.Parser.Parser().parsestr(fulltext)       # may fail!
        except:
            return self.errorMessage     # or let call handle? can check return

    def parseMessageRaw(self, fulltext):
        """
        parse headers only, return root email.Message object
        stops after headers parsed, for efficiency (not yet used here)
        payload of message object is raw text of mail after headers
        """
        try:
            return email.Parser.HeaderParser().parsestr(fulltext)
        except:
            return self.errorMessage


####################################
# send messages, add attachments
####################################


class MailSender(MailTool):
    """
    send mail: format message, interface with SMTP server
    works on any machine with Python+Inet, doesn't use cmdline mail
    """
    def __init__(self, smtpserver=None):
        self.smtpServerName  = smtpserver or mailconfig.smtpservername

    def sendMessage(self, From, To, Subj, extrahdrs, bodytext, attaches):
        """
        format,send mail: blocks caller, thread me in a gui
        bodytext is main text part, attaches is list of filenames
        extrahdrs is list of (name, value) tuples to be added
        raises uncaught exception if send fails for any reason
        
        assumes that To, Cc, Bcc hdr values are lists of 1 or more already
        stripped addresses (possibly in full name+<addr> format); client
        must split these on delimiters, parse, or use multi-line input;
        note that smtp allows full name+<addr> format in recipients
        """
        if not attaches:
            msg = Message()
            msg.set_payload(bodytext)
        else:
            msg = MIMEMultipart()
            self.addAttachments(msg, bodytext, attaches)

        recip = To
        msg['From']    = From
        msg['To']      = ', '.join(To)              # poss many: addr list
        msg['Subject'] = Subj                       # servers reject ';' sept
        msg['Date']    = email.Utils.formatdate()   # curr datetime, rfc2822 utc
        for name, value in extrahdrs:               # Cc, Bcc, X-Mailer, etc.
            if value:
                if name.lower() not in ['cc', 'bcc']:
                    msg[name] = value
                else:
                    msg[name] = ', '.join(value)    # add commas between
                    recip += value                  # some servers reject ['']
        fullText = msg.as_string()                  # generate formatted msg

        # sendmail call raises except if all Tos failed,
        # or returns failed Tos dict for any that failed
        
        self.trace('Sending to...'+ str(recip))
        self.trace(fullText[:256])
        server = smtplib.SMTP(self.smtpServerName)           # this may fail too
        try:
            failed = server.sendmail(From, recip, fullText)  # except or dict
        finally:
            server.quit()                                    # iff connect okay
        if failed:
            class SomeAddrsFailed(Exception): pass
            raise SomeAddrsFailed('Failed addrs:%s\n' % failed)
        self.trace('Send exit')

    def addAttachments(self, mainmsg, bodytext, attaches):
        # format a multi-part message with attachments
        msg = MIMEText(bodytext)                 # add main text/plain part
        mainmsg.attach(msg)
        for filename in attaches:                # absolute or relative paths
            if not os.path.isfile(filename):     # skip dirs, etc.
                continue
            
            # guess content type from file extension, ignore encoding
            contype, encoding = mimetypes.guess_type(filename)
            if contype is None or encoding is not None:  # no guess, compressed?
                contype = 'application/octet-stream'     # use generic default
            self.trace('Adding ' + contype)

            # build sub-Message of apropriate kind
            maintype, subtype = contype.split('/', 1)
            if maintype == 'text':
                data = open(filename, 'r')
                msg  = MIMEText(data.read(), _subtype=subtype)
                data.close()
            elif maintype == 'image':
                data = open(filename, 'rb')
                msg  = MIMEImage(data.read(), _subtype=subtype)
                data.close()
            elif maintype == 'audio':
                data = open(filename, 'rb')
                msg  = MIMEAudio(data.read(), _subtype=subtype)
                data.close()
            else:
                data = open(filename, 'rb')
                msg  = MIMEBase(maintype, subtype)
                msg.set_payload(data.read())
                data.close()                            # make generic type
                email.Encoders.encode_base64(msg)       # encode using Base64

            # set filename and attach to container
            basename = os.path.basename(filename)
            msg.add_header('Content-Disposition',
                           'attachment', filename=basename)
            mainmsg.attach(msg)

        # text outside mime structure, seen by non-MIME mail readers
        mainmsg.preamble = 'A multi-part MIME format message.\n'
        mainmsg.epilogue = ''  # make sure message ends with a newline


####################################
# retrieve mail from a pop server
####################################


class MailFetcher(MailTool):
    """
    fetch mail: connect, fetch headers+mails, delete mails
    works on any machine with Python+Inet; subclass me to cache
    implemented with the POP protocol; IMAP requires new class
    """
    def __init__(self, popserver=None, popuser=None, hastop=True):
        self.popServer   = popserver or mailconfig.popservername
        self.popUser     = popuser   or mailconfig.popusername
        self.srvrHasTop  = hastop
        self.popPassword = None
                 
    def connect(self):
        self.trace('Connecting...')
        self.getPassword()                          # file, gui, or console
        server = poplib.POP3(self.popServer)  
        server.user(self.popUser)                   # connect,login pop server
        server.pass_(self.popPassword)              # pass is a reserved word
        self.trace(server.getwelcome())             # print returned greeting 
        return server

    def downloadMessage(self, msgnum):
        # load full text of one msg
        self.trace('load '+str(msgnum))
        server = self.connect()
        try:
            resp, msglines, respsz = server.retr(msgnum)
        finally:
            server.quit()
        return '\n'.join(msglines)                 # concat lines for parsing

    def downloadAllHeaders(self, progress=None, loadfrom=1):
        """
        get sizes, headers only, for all or new msgs
        begins loading from message number loadfrom
        use loadfrom to load newly-arrived mails only
        use downloadMessage to get the full msg text
        """
        if not self.srvrHasTop:                    # not all servers support TOP
            return self.downloadAllMsgs(progress)  # naively load full msg text
        else:
            self.trace('loading headers')
            server = self.connect()                # mbox now locked until quit
            try:
                resp, msginfos, respsz = server.list()   # 'num size' lines list
                msgCount = len(msginfos)                 # alt to srvr.stat[0]
                msginfos = msginfos[loadfrom-1:]         # drop already-loadeds
                allsizes = [int(x.split()[1]) for x in msginfos]
                allhdrs  = []
                for msgnum in range(loadfrom, msgCount+1):          # poss empty
                    if progress: progress(msgnum, msgCount)         # callback?
                    resp, hdrlines, respsz = server.top(msgnum, 0)  # hdrs only
                    allhdrs.append('\n'.join(hdrlines))
            finally:
                server.quit()                          # make sure unlock mbox
            assert len(allhdrs) == len(allsizes)
            self.trace('load headers exit')
            return allhdrs, allsizes, False

    def downloadAllMessages(self, progress=None, loadfrom=1):
        # load full message text for all msgs, despite any caching
        self.trace('loading full messages')
        server = self.connect()
        try:
            (msgCount, msgBytes) = server.stat()
            allmsgs  = []
            allsizes = []
            for i in range(loadfrom, msgCount+1):         # empty if low >= high
                if progress: progress(i, msgCount)
                (resp, message, respsz) = server.retr(i)  # save text on list
                allmsgs.append('\n'.join(message))        # leave mail on server
                allsizes.append(respsz)                   # diff from len(msg)
        finally:
            server.quit()                                 # unlock the mail box
        assert len(allmsgs) == (msgCount - loadfrom) + 1  # msg nums start at 1
        assert sum(allsizes) == msgBytes
        return allmsgs, allsizes, True

    def deleteMessages(self, msgnums, progress=None):
        # delete multiple msgs off server
        server = self.connect()
        try:
            for (ix, msgnum) in enumerate(msgnums):   # dont reconnect for each
                if progress: progress(ix+1, len(msgnums))
                server.dele(msgnum)                     
        finally:                                      # changes msgnums: reload
            server.quit()
        
    def getPassword(self):
        """
        get pop password if not yet known
        not required until go to server
        from client-side file or subclass method
        """
        if not self.popPassword:
            try:
                localfile = open(mailconfig.poppasswdfile)
                self.popPassword = localfile.readline()[:-1]
                self.trace('local file password' + repr(self.popPassword))
            except:
                self.popPassword = self.askPopPassword()

    def askPopPassword(self):
        assert False, 'Subclass must define method'


class MailFetcherConsole(MailFetcher):
    def askPopPassword(self):
        import getpass
        prompt = 'Password for %s on %s?' % (self.popUser, self.popServer)
        return getpass.getpass(prompt)


##################################
# self-test when run as a program
##################################


if __name__ == '__main__':
    sender = MailSender()
    sender.sendMessage(From      = mailconfig.myaddress,
                       To        = [mailconfig.myaddress],
                       Subj      = 'testing 123',
                       extrahdrs = [('X-Mailer', 'mailtools.py')],
                       bodytext  = 'Here is my source code',
                       attaches  = ['mailtools.py'])

    fetcher = MailFetcherConsole()
    def status(*args): print args
    
    hdrs, sizes, all = fetcher.downloadAllHeaders(status)
    for num, hdr in enumerate(hdrs[:5]):
        print hdr
        if raw_input('load mail?') in ['y', 'Y']:
            print fetcher.downloadMessage(num+1), '\n', '-'*70
            
    msgs, sizes, all = fetcher.downloadAllMessages(status)
    for msg in msgs[:5]:
        print msg, '\n', '-'*70
    raw_input('Press Enter to exit')