'''
mimeparse module, version 0.4.4

(c) 2003-2004 Jan ONDREJ (SAL) <ondrejj(at)salstar.sk>
                                                                                
 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.
'''

from avlib import *
import os,re,filetype
import email,mimetypes,email.Errors,email.Parser
import email.Utils,binascii,uu,base64
import multifile,mimetools
from StringIO import StringIO

__all__=['mimeparse','mime_ext']

class mime_ext_int(ascanner):
      '''
      Internal scanner for mime_ext.
      
      This scanner checks defined extension of file.
      Regexp check with ignoring of case is used.
      '''
      def __init__(self,wrong_type):
          self.filename=''
          self.wrong_type=wrong_type
      def scanbuffer(self,data,args={}):
          if (self.wrong_type!='') and (type(self.filename)==type('string')):
            if re.search(self.wrong_type,self.filename,re.IGNORECASE):
               return 1.0,'MatchedExtension', [self.filename+' matched '+self.wrong_type+'\n']
          return 0.0,'',[]

def decode_part(self,filename=''):
    '''
    Decode a part of message. Return at least some bytes from this part,
    if error occures.
    '''
    if self.is_multipart():
      return None,'multipart'
    cte = self.get('content-transfer-encoding', '').lower()
    debug.echo(5,"mimeparse(): decoding: ",filename,", type: ",cte)
    payload=self.get_payload()
    if cte == 'quoted-printable':
      return email.Utils._qdecode(payload),''
    elif cte == 'base64':
      try:
        return binascii.a2b_base64(payload),''
      except binascii.Error:
        # try another method to decode base64 buffer
        sfp = StringIO()
        try:
          base64.decode(StringIO(payload+'\n'), sfp)
          return sfp.getvalue(),''
        except:
          debug.echo(2,"ERROR: mimeparse(): part_decode: Incorrect padding")
          return sfp.getvalue(),'truncated'
    elif cte in ('x-uuencode', 'uuencode', 'uue', 'x-uue'):
      sfp = StringIO()
      try:
        uu.decode(StringIO(payload+'\n'), sfp)
        return sfp.getvalue(),''
      except uu.Error:
        debug.echo(2,"ERROR: mimeparse(): part_decode: Some decoding problem")
        return sfp.getvalue(),'truncated'
    return payload,'encoded'

def decode_email_alt(data,scanner,args):
    '''
    Alternate method to decode an email.
    '''
    try:
      buffer=StringIO(data)
      msg = mimetools.Message(buffer)
      msgtype = msg.gettype()
      params = msg.getplist()
    except:
      debug.echo(1, "MimeParse: ERROR: mimetools.Message(), all decoders failed")
      debug.traceback(6, "MimeParse: ")
      return '',[],[]
    files=[]
    counter=0
    level,vir,ret=0.0,'',[]
    if msgtype[:10] == "multipart/":
      file = multifile.MultiFile(buffer)
      file.push(msg.getparam("boundary"))
      while file.next():
        counter+=1
        submsg = mimetools.Message(file)
        sfn=submsg.getparam('name')
        if sfn:
          scanner.filename=normalize_filename(sfn)
        else:
          try:
            ext = mimetypes.guess_extension(submsg.gettype())
            if not ext:
              ext=".bin"
          except:
            ext=".bin"
          scanner.filename = 'part-%03d%s' % (counter,ext)
        debug.echo(5,"mimeparse(): decoding: ",scanner.filename)
        try:
          data = StringIO()
          mimetools.decode(file, data, submsg.getencoding())
        except binascii.Error:
          pass
        buffer=data.getvalue()
        if buffer:
          debug.echo(6,'mimeparse(): decoded len=',len(buffer))
          if scanner.is_multibuffer:
            scanner.addbuffer(buffer)
          else:
            l,v,r=scanner.scanbuffer(buffer,args)
            if not vir:
              vir=v
            ret+=r
            level+=l
          files.append(scanner.filename)
      file.pop()
    if scanner.is_multibuffer:
      level,vir,ret=scanner.scanmultibuffer([],args)
    del scanner
    return level,vir,ret,files

def decode_email(data,scanner,args):
    '''
    Decode whole email and run a scanner on decoded parts.
    '''
    try:
      msg=email.message_from_string(data)
    except email.Errors.BoundaryError,es:
      debug.echo(2, "MimeParse: ERROR: BoundaryError: ",es)
      return decode_email_alt(data,scanner,args)
    except email.Errors.HeaderParseError,es:
      debug.echo(2, "MimeParse: ERROR: HeaderParseError: ",es)
      return decode_email_alt(data,scanner,args)
    files=[]
    counter=1
    level,vir,ret=0.0,'',[]
    for part in msg.walk():
      counter+=1
      try:
        sfn=part.get_filename()
        if sfn:
          scanner.filename=normalize_filename(sfn)
        else:
          # no filename, make own
          try:
            ext = mimetypes.guess_extension(part.get_type())
          except:
            ext='.bin'
          if not ext:
            ext='.bin'
          scanner.filename = 'part-%03d%s' % (counter, ext)
        buffer,dstat=decode_part(part,scanner.filename)
        if buffer!=None:
          debug.echo(5,"mimeparse(): buffer len=",len(buffer)," ",dstat)
      except:
        buffer=""
      if buffer:
        if scanner.is_multibuffer:
          scanner.addbuffer(buffer)
        else:
          l,v,r=scanner.scanbuffer(buffer,args)
          if not vir:
            vir=v
          ret+=r
          level+=l
        files.append(scanner.filename)
    if scanner.is_multibuffer:
      level,vir,ret=scanner.scanmultibuffer([],args)
    del scanner
    return level,vir,ret,files

class mimeparse(interscanner):
      '''
      Email parser interscanner. [obsolete]
      
      This scanner parses MIME attached mail and send multiple of
      buffers into another_scanner(). If an multiscanner is used
      as scanner1(), optimal method is used (scans all files
      at once).

      Usage: mimeparse(scanner1(...))
      
      This scanner is obsolete since 0.5.0, use parsemail() instead.
      '''
      name='MimeParse()'
      def __init__(self,scanner):
          self.name="MimeParse("+scanner.name+")"
          self.scanner=scanner
      def scanbuffer(self,buffer,args={}):
          self.scanner.prescan()
          level,vir,ret,filelist=decode_email(buffer,self.scanner,args)
          self.scanner.postscan(level,vir,ret)
          return level,vir,ret

class mime_ext(ascanner):
      '''
      A extension scanner. [obsolete]
      
      This scanner checks defined extension of file in mime attachments.
      Regexp checks with ignoring of case are used.
      You can use predefined extensions from mime_ext.EXECUTABLES.

      Usage: mime_ext(wrong_type='')
      
      Example: mime_ext(mime_ext.EXECUTABLES)
      
      This scanner is obsolete since 0.5.9, use attach_name() instead.
      '''
      name='MimeExt()'
      EXECUTABLES='\.(exe|com|pif|lnk|scr|vbs|cpl)$'
      def __init__(self,wrong_type=''):
          self.wrong_type=wrong_type
      def scanbuffer(self,buffer,args={}):
          level,vir,ret,filelist=decode_email(
            buffer,mime_ext_int(self.wrong_type),args)
          if vir:
            globals.found_by=self
          return level,vir,ret
