'''
mimeparse module, version 0.4.3

(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','file_type']

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 scanstream(self,data):
          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 stream
        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):
    '''
    Alternate method to decode an email.
    '''
    try:
      stream=StringIO(data)
      msg = mimetools.Message(stream)
      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(stream)
      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
        stream=data.getvalue()
        if stream:
          debug.echo(6,'mimeparse(): decoded len=',len(stream))
          if scanner.is_multistream:
            scanner.addstream(stream)
          else:
            l,v,r=scanner.scanstream(stream)
            if not vir:
              vir=v
            ret+=r
            level+=l
          files.append(scanner.filename)
      file.pop()
    if scanner.is_multistream:
      level,vir,ret=scanner.scanmultistream([])
    del scanner
    return level,vir,ret,files

def decode_email(data,scanner):
    '''
    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)
    except email.Errors.HeaderParseError,es:
      debug.echo(2, "MimeParse: ERROR: HeaderParseError: ",es)
      return decode_email_alt(data,scanner)
    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)
        stream,dstat=decode_part(part,scanner.filename)
        if stream!=None:
          debug.echo(5,"mimeparse(): stream len=",len(stream)," ",dstat)
      except:
        stream=""
      if stream:
        if scanner.is_multistream:
          scanner.addstream(stream)
        else:
          l,v,r=scanner.scanstream(stream)
          if not vir:
            vir=v
          ret+=r
          level+=l
        files.append(scanner.filename)
    if scanner.is_multistream:
      level,vir,ret=scanner.scanmultistream([])
    del scanner
    return level,vir,ret,files

class mimeparse(interscanner):
      '''
      mimeparse interscanner
      
      Usage: mimeparse(another_scanner(...))
      
      This scanner parses MIME attached mail and send multiple of
      streams into another_scanner(). If an multiscanner is used
      as another_scanner, optimal method is used (scans all files
      at once).
      '''
      name='MimeParse()'
      def __init__(self,scanner):
          self.name="MimeParse("+scanner.name+")"
          self.scanner=scanner
      def scanstream(self,stream):
          self.scanner.prescan()
          level,vir,ret,filelist=decode_email(stream,self.scanner)
          self.scanner.postscan(level,vir,ret)
          return level,vir,ret

class mime_ext(ascanner):
      '''
      mime_ext() is a extension scanner.

      Usage: mime_ext('wrong types regular expressios')
      
      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,
      for example:
        mime_ext(mime_ext.EXECUTABLES)
      '''
      name='MimeExt()'
      EXECUTABLES='(\.exe|\.com|\.pif|\.lnk|\.scr|\.vbs)$'
      def __init__(self,wrong_type=''):
          self.wrong_type=wrong_type
      def scanstream(self,stream):
          level,vir,ret,filelist=decode_email(stream,mime_ext_int(self.wrong_type))
          if vir:
            globals.found_by=self
          return level,vir,ret

class file_type(ascanner):
      '''
      file_type() is a realscanner, which scans type of a file
      
      Usage: file_type({'type':'vir',...})
      
      Where 'type' is a name of type returned by filetype.* function
            'vir' is a vir name returned, if file type is matched
      '''
      name='file_type()'
      def __init__(self,types):
          self.types=types
      def scanstream(self,stream):
          fta=filetype.what(stream)
          for ft in fta:
            if ft in self.types:
              return 1.0,self.types[ft],[self.filename+": "+str(fta)]
          return 0.0,'',[]
      def scanfile(self,files,dirname=''):
          for fname in files:
            fta=filetype.file(fname)
            for ft in fta:
              if ft in self.types:
                return 1.0,self.types[ft],[fname+": "+str(fta)]
          return 0.0,'',[]
