'''
Header interscanners for sagator, version 0.6.2

(c) 2003-2019 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 __future__ import absolute_import

from avlib import *
import re
from .match import match_any

__all__ = ['add_header', 'modify_header', 'modify_subject', 'remove_headers']

class add_header(match_any):
  '''
  Add email headers.

  This scanner can be used to add an header line to email.

  Usage: add_header('Header_name','header_data',scanners)
     or: add_header(...).onlyif('virname')
  
  Where: 'Header_name' is a string, which defines name of the header
         'header_data' is a string, which defines header's content
         'virname' is a string, which defines any virus name.
           By default all virus names at least 1 character long: '.'
         There are some special pattern, which are replaced:
           %V ... virus name
           %S ... scanner name
           %L ... detected level
  '''
  name = 'add_header()'
  def __init__(self, hdr, data, *scanners):
      self.ONLY_IF = b'.'
      self.HDR = hdr
      self.DATA = data
      match_any.__init__(self, scanners)
  def onlyif(self, pattern=b'.'):
      self.ONLY_IF = pattern
      return self
  def hdr_replace(self, to, level, detected):
      return re.sub("%V", tostr(detected),
               re.sub("%S", globals.found_by.name,
                 re.sub("%L", str(level),
                   to)))
  def scanbuffer(self, buffer, args={}):
      level, detected, virlist = match_any.scanbuffer(self, buffer, args)
      if re.search(self.ONLY_IF, detected, re.IGNORECASE):
        mail.addheader(self.HDR, self.hdr_replace(self.DATA, level, detected))
      return level, detected, virlist

class modify_header(add_header):
  '''
  Modify email headers.

  This scanner can be used to modify email headers.

  Usage: modify_header('^Header_regexp:','Header_regexp: prefix',scanners)
     or: modify_header(...).onlyif('virname')
  
  Where: '^Header_regexp:' is a string, which defines regular expression
           to find for this header
         'Header_regexp: prefix' is a string, which defines new header
           data
         'virname' is a string, which defines any virus name.
           By default all virus names at least 1 character long: '.'
         There are some special pattern, which are replaced:
           %V ... virus name
           %S ... scanner name
           %L ... detected level
  '''
  name = 'modify_header()'
  def __init__(self, frm, to, *scanners):
      self.ONLY_IF = b'.'
      self.MOD_FRM = tobytes(frm)
      self.MOD_TO = tobytes(to)
      match_any.__init__(self,scanners)
  def hdr_replace(self, to, level, detected):
      return re.sub(b"%V", detected,
               re.sub(b"%S", tobytes(globals.found_by.name),
                 re.sub(b"%L", tobytes(str(level)),
                   to)))
  def scanbuffer(self, buffer, args={}):
      level, detected, virlist = match_any.scanbuffer(self, buffer, args)
      if re.search(self.ONLY_IF, detected, re.IGNORECASE):
        change_to = self.hdr_replace(self.MOD_TO, level, detected)
        changes = mail.modheader(self.MOD_FRM, change_to)
        if changes==0:
          mail.addheader(change_to)
      return level, detected, virlist

class modify_subject(modify_header):
  '''
  Modify email Subject.

  This scanner can be used to modify email Subject.

  Usage: modify_subject('prefix',scanners)
     or: modify_subject(...).onlyif('virname')
  
  Where: 'prefix' is an prefix added into Subject.
         'virname' is a string, which defines any virus name.
           By default all virus names at least 1 character long: '.'
         For more information see modify_header() interscanner.

  For example original Subject line is:
    Subject: Testing mail
  and i use:
    modify_subject('[SPAM]', spamassassind(['localhost',783]))
  When spam will be found, header will be modifyed to:
    Subject: [SPAM] Testing mail
  '''
  name = 'modify_subject()'
  def __init__(self, prefix, *scanners):
      modify_header.__init__(self, "^Subject:", "Subject: "+prefix, scanners)

class remove_headers(ascanner):
  '''
  Remove email headers.

  This scanner can be used to remove email headers.

  Usage: remove_headers('header name' [,'next header name' [, ...] ])
  
  Where: 'header name' is a string, which defines header key to delete
  '''
  name = 'remove_headers()'
  def __init__(self, *hdrs):
      self.DEL_HDRS = hdrs
  def scanbuffer(self, buffer, args={}):
      for hdr in self.DEL_HDRS:
        mail.delheader(tobytes(hdr))
      return 0.0, b'', []