#!/usr/bin/env python

#
# Copyright 2015, Olivier MATZ <zer0@droids-corp.org>
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the University of California, Berkeley nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE REGENTS AND CONTRIBUTORS BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

import imapami.utils

import copy
from email.mime.multipart import MIMEMultipart
from email.mime.message import MIMEMessage
from email.utils import formatdate
import logging
import shlex
import smtplib
import subprocess

_all_actions = {}
_LOGLEVELS = [logging.CRITICAL, logging.ERROR, logging.WARNING,
              logging.INFO, logging.DEBUG]

def register(action_class):
    """
    Register an action class
    """
    if not issubclass(action_class, ImapamiAction):
        raise ValueError('"%s" action is not a subclass of ImapamiAction' %
                         action_class)
    if action_class.name in _all_actions:
        raise ValueError('"%s" action is already registered' %
                         action_class.name)
    _all_actions[action_class.name] = action_class

class ImapamiAction(object):
    """
    This is the parent class for actions.

    An action is a task executed by imapami for a mail when
    the conditions match. For instance, it can move the mail in
    another directory, delete a mail, set a variable, execute a
    script...
    """
    name = None

    def __init__(self, fetch=None, terminal=False):
        """
        Generic action constructor.
        :arg str fetch:
          Set to "headers", "part1", "all" if it is needed to fetch
          the header, the body or the attachments of the mail
        :arg boolean terminal:
          Set to true if the action has to be the last of the list
          (ex: the mail is deleted or moved)
        """
        self.fetch = fetch or "no"
        self.terminal = terminal

    def evaluate(self, arg, ami, hdrs):
        """
        Evaluate a string argument

        Replace variables and headers in arg by their values and
        return it.

        :arg string arg:
          The argument to be evaluated
        :arg Imapami ami:
          The imapami object, containing the variables
        :arg email.message.Message hdrs:
          The headers of the mail, or None if not available
        :returns:
          The evaluated argument.
        """
        if hdrs is not None:
            variables = imapami.utils.headers_to_unicode(hdrs)
        else:
            variables = {}
        variables.update(ami.variables)
        fmt = imapami.utils.VarFormatter()
        arg = fmt.format(unicode(arg), **variables)
        return arg

    def process(self, imap, mail):
        """
        Process the action

        :arg Imapami ami:
          The Imapami object
        :arg ImapamiMail mail:
          The mail data
        :returns:
          True on success, False on error
        """
        return True

class ImapamiActionList(ImapamiAction):
    """
    Execute a list of actions.

    This command is used internally when the list of actions of a rule
    is parsed. There is no real need for a user, as the rule already contains
    a list of actions.

    Arguments:
      List of actions

    Example:
      list:
      - copy: {dest: archive}
      - move: {dest: mbox}
    """
    name = "list"
    def __init__(self, *action_list):
        action_list = [new(a) for a in action_list]
        for a in action_list[:-1]:
            if a.terminal == True:
                raise ValueError("terminal action in the middle of an action list")
        fetch = imapami.utils.highest_fetch_level(
            [a.fetch for a in action_list])
        terminal = action_list[-1].terminal
        ImapamiAction.__init__(self, fetch, terminal)
        self.action_list = action_list
    def process(self, ami, mail):
        ret = True
        for a in self.action_list:
            if a.process(ami, mail) == False:
                ret = False
        return ret
register(ImapamiActionList)

class ImapamiActionCopy(ImapamiAction):
    """
    Copy the mail in another directory.

    Arguments:
      dest: <string>
        destination directory

    Example:
      copy: {dest: my_subdir}
    """
    name = "copy"
    def __init__(self, dest):
        ImapamiAction.__init__(self)
        self.dest = dest
    def process(self, ami, mail):
        imap = ami.imap
        dest = self.evaluate(self.dest, ami, mail.msg)
        imap.create(dest)
        ret, msg = imap.uid("COPY", mail.item, dest)
        if ret != "OK":
            ami.logger.warning(
                "imap copy returned %s: %s" % (ret, str(msg)))
            return False
        return True
register(ImapamiActionCopy)

class ImapamiActionChangeFlag(ImapamiAction):
    """
    Change an IMAP flag of a mail.

    Arguments:
      flag: <string>
        Name of the IMAP flag (ex: 'Seen', 'Answered', 'Deleted', ...)
        Refer to RFC3501 for details.
      enable: <boolean>
        True to set the flag, False to reset it.

    Example:
      change-flag: {flag: Seen, enable: True}
    """
    name = "change-flag"
    def __init__(self, flag, enable):
        ImapamiAction.__init__(self)
        self.enable = enable
        self.flag = flag
    def process(self, ami, mail):
        imap = ami.imap
        if self.enable == True:
            cmd = '+FLAGS'
        else:
            cmd = '-FLAGS'
        flag = '(\\%s)' % self.evaluate(self.flag, ami, mail.msg)
        ret, msg = imap.uid("STORE", mail.item, cmd, flag)
        if ret != "OK":
            ami.logger.warning(
                "imap store '%s %s' returned %s: %s" % (
                    cmd, flag, ret, str(msg)))
            return False
        return True
register(ImapamiActionChangeFlag)

class ImapamiActionDelete(ImapamiActionChangeFlag):
    """
    Mark a mail as deleted.

    Arguments: None

    Example:
      deleted: {}
    """
    name = "deleted"
    def __init__(self):
        ImapamiActionChangeFlag.__init__(self, "Deleted", True)
        self.terminal = True
register(ImapamiActionDelete)

class ImapamiActionSeen(ImapamiActionChangeFlag):
    """
    Mark a mail as seen.

    Arguments: None

    Example:
      seen: {}
    """
    name = "seen"
    def __init__(self):
        ImapamiActionChangeFlag.__init__(self, "Seen", True)
register(ImapamiActionSeen)

class ImapamiActionUnseen(ImapamiActionChangeFlag):
    """
    Mark a mail as not seen.

    Arguments: None

    Example:
      unseen: {}
    """
    name = "unseen"
    def __init__(self):
        ImapamiActionChangeFlag.__init__(self, "Seen", False)
register(ImapamiActionUnseen)

class ImapamiActionMove(ImapamiAction):
    """
    Move the mail in another directory.

    Arguments:
      dest: <string>
        destination directory

    Example:
      move: {dest: my_subdir}
    """
    name = "move"
    def __init__(self, dest):
        ImapamiAction.__init__(self, terminal=True)
        self.dest = dest
    def process(self, ami, mail):
        imap = ami.imap
        dest = self.evaluate(self.dest, ami, mail.msg)
        imap.create(dest)
        ret, msg = imap.uid("COPY", mail.item, dest)
        if ret != "OK":
            ami.logger.warning(
                "imap copy returned %s: %s" % (ret, str(msg)))
            return False
        ret, msg = imap.uid("STORE", mail.item, '+FLAGS', '(\\Deleted)')
        if ret != "OK":
            ami.logger.warning(
                "imap delete returned %s: %s" % (ret, str(msg)))
            return False
        return True
register(ImapamiActionMove)

class ImapamiActionLog(ImapamiAction):
    """
    Log a message.

    Arguments:
      msg: <string>
        Message to be logged. Like all action string arguments, it
        can contain variables or header name.
      level: <integer> [optional, default = 3 (info)]
        The level of the log from 0 (critical) to 4 (debug). Note that
        it must be lower or equal to the global log level, else it
        won't appear in the log file.

    Example:
      log: {msg: 'Mail from {From} is received', level: 3}
    """
    name = "log"
    def __init__(self, msg, level=3):
        ImapamiAction.__init__(self, fetch="headers")
        self.msg = msg
        if not isinstance(level, int) or level < 0 or level > 4:
            raise ValueError("invalid log level")
        self.level = level
    def process(self, ami, mail):
        msg = self.evaluate(self.msg, ami, mail.msg)
        l = _LOGLEVELS[self.level]
        ami.logger.log(l, msg)
        return True
register(ImapamiActionLog)

class ImapamiActionPipe(ImapamiAction):
    """
    Pipe the mail to an external program.

    Arguments:
      command: <string>
        The command to run.
      what: <string> [default='headers']
        Define what is sent to the program. Valid values are:
          'headers': send the headers
          'body_part1': send the first part of the message
          'body_all': send all the message body
          'headers+body_part1': send headers and first part of the message
          'headers+body_all': send all
      shell: <boolean> [default=False]
        Invoke a shell for this command. This allows for instance to use
        a pipe or a redirection to a file.

    Example:
      pipe: {command: 'cat > /tmp/foobar', shell: True}
    """
    name = "pipe"
    def __init__(self, command, what='headers', shell=False):
        valid = ['headers', 'body_part1', 'body_all',
                 'headers+body_part1', 'headers+body_all']
        if not what in valid:
            raise ValueError("invalid 'what' field. Valid values are: %s" %
                             valid)
        if what in ['headers']:
            fetch = "headers"
        elif what in ['body_part1', 'headers+body_part1']:
            fetch = "part1"
        else:
            fetch = "all"
        ImapamiAction.__init__(self, fetch=fetch)
        self.command = command
        self.what = what
        self.shell = shell
    def process(self, ami, mail):
        if self.shell == True:
            command = self.command
        else:
            command = shlex.split(self.command)
        process = subprocess.Popen(command,
                                   stdin=subprocess.PIPE,
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE,
                                   shell=self.shell)
        data = ''
        if self.what in ['headers', 'headers+body_part1', 'headers+body_all']:
            data += str(mail.msg)
        if self.what in ['body_part1', 'headers+body_part1']:
            data += mail.body_part1
        if self.what in ['body_all', 'headers+body_all']:
            data += mail.body_all
        stdout, stderr = process.communicate(data)
        ret = process.wait()
        ami.variables["stdout"] = stdout
        ami.variables["stderr"] = stderr
        ami.variables["ret"] = ret
        return True
register(ImapamiActionPipe)

class ImapamiActionSetHeader(ImapamiAction):
    """
    Set or remove header fields in the mail.

    Dpending on the mode, it is possible to:
    - append it at the end of the mail headers (mode=add).
    - remove all occurences of a header (mode=del)
    - replace the first occurence of a header (mode=replace)

    As IMAP does not allow to modify a mail, this action creates
    a modified copy of the original mail, and delete the original
    one.

    Arguments:
      <list of commands (see below)>

    Format of a command:
      mode: <string>
        Determine if headers should be added, removed, replaced. Can
        be 'add', 'del', 'replace'.
      <key (string)>: <value (string)>
        The key contains the name of the header field (ex: Subject),
        associated to the value to be set (None for delete operations).
        Multiple fields can be specified.

    Example:
      set-header: [{mode: del, Message-Id: None},
                   {mode: add, Foo: bar, X-mailer: toto},
                   {mode: replace, Subject: 'subject was: {Subject}'}]
    """
    name = "set-header"
    def __init__(self, *args):
        ImapamiAction.__init__(self, fetch="all", terminal=True)
        for cmd in args:
            if not isinstance(cmd, dict):
                raise ValueError("set-header command is not a dict")
            if not cmd.has_key("mode"):
                raise ValueError("set-header command has no mode")
            if cmd["mode"] not in ['add', 'del', 'replace']:
                raise ValueError("invalid mode for set-header command")
        self.cmdlist = args

    def process(self, ami, mail):
        parsed_headers = copy.deepcopy(mail.msghdrs)
        for cmd in self.cmdlist:
            mode = cmd.pop("mode")
            fields = cmd
            if mode == "add":
                for k, v in fields.iteritems():
                    parsed_headers[k] = v
            elif mode == "del":
                for k, v in fields.iteritems():
                    del parsed_headers[k]
            elif mode == "replace":
                for k, v in fields.iteritems():
                    parsed_headers.replace_header(k, v)

        ami.imap.append(mail.inbox, mail.flags, mail.internal_date,
                        str(parsed_headers) + mail.body_all)
        return True
register(ImapamiActionSetHeader)

class ImapamiSmtpSender(object):
    """
    A SMTP sender using smtp lib
    """

    def __init__(self, host, port=None, encryption=None, login=None,
                 password=None):
        """
        Initialize the SMTP sender class

        :arg string host:
          Hostname or ip address of the smtp server.
        :arg integer port:
          The port of the smtp server to connect to.
        :arg string encryption:
          Define encryption type for smtp: "none", "ssl", "starttls".
          Default is "none".
        :arg string login:
          User login for smtp server. If no login is specified, assume no
          login is required.
        :arg string password:
          Password associated to the login.
        """
        self.host = host
        self.port = port
        self.encryption = encryption or "none"
        self.login = login
        self.password = password

    def send(self, ami, sender, to, mail):
        """
        Send a mail.

        :arg Imapami ami:
          The Imapami object
        :arg string mail:
          The mail to send
        :arg string sender:
          The sender of the mail
        :arg string to:
          To mail address
        """
        ami.logger.debug("sending to smtp server %s" % self.host)
        args = [self.host]
        if self.port is not None:
            args.append(self.port)
        if self.encryption == "ssl":
            smtp = smtplib.SMTP_SSL(*args)
        else:
            smtp = smtplib.SMTP(*args)
        try:
            if self.encryption == "starttls":
                smtp.starttls()
            if self.login is not None:
                smtp.login(self.loging, self.password)
            smtp.sendmail(sender, to, mail)
            smtp.close()
        except smtplib.SMTPAuthenticationError as e:
            ami.logger.warning("smtp authentication error: %s", str(e))
            return False
        except (smtplib.SMTPException, smtplib.SMTPHeloError,
                smtplib.SMTPRecipientsRefused, smtplib.SMTPSenderRefused,
                smtplib.SMTPDataError, RuntimeError) as e:
            ami.logger.warning("smtp error: %s", str(e))
            return False
        return True

class ImapamiActionForward(ImapamiAction):
    """
    Forward a mail.

    The mail is forwarded with its attachments. New mail headers
    are created from scratch. The subject of the mail is prefixed
    with 'Fwd: '.

    Arguments:
      sender: <string>
        The sender of the mail.
      to: <string>
        The of the forwarded mail.
      host: <string> (mandatory)
        Hostname or ip address of the smtp server.
      port: <integer> (optional)
        The port of the smtp server to connect to.
      encryption: <string> (optional)
        Define encryption type for smtp: "none", "ssl", "starttls".
        Default is none.
      login: <string> (optional)
        User login for smtp server. If no login is specified, assume no
        login is required.
      password: <string> (optional)
        Password associated to the login.

    Example:
      forward: {sender: foo@example.com, to: toto@example.com,
                host: mail.example.com}
    """
    name = "forward"
    def __init__(self, sender, to, host, port=None, encryption=None, login=None,
                 password=None):
        ImapamiAction.__init__(self, fetch="all")
        self.sender = sender
        self.to = to
        self.smtp = ImapamiSmtpSender(host, port, encryption, login, password)

    def process(self, ami, mail):
        sender = self.evaluate(self.sender, ami, mail.msg)
        to = self.evaluate(self.to, ami, mail.msg)
        subject = "Fwd: " + self.evaluate('{Subject}', ami, mail.msg)
        new_mail = MIMEMultipart()
        new_mail['Date'] = formatdate(localtime=True)
        new_mail['From'] = sender
        new_mail['To'] = to
        new_mail['Subject'] = subject
        new_mail.attach(MIMEMessage(mail.msg))
        return self.smtp.send(ami, sender, to, new_mail.as_string())
register(ImapamiActionForward)

class ImapamiActionBounce(ImapamiAction):
    """
    Bounce a mail.

    The mail is transfered with its attachments without modification.

    Arguments:
      sender: <string>
        The sender of the mail.
      to: <string>
        The of the bounceed mail.
      host: <string> (mandatory)
        Hostname or ip address of the smtp server.
      port: <integer> (optional)
        The port of the smtp server to connect to.
      encryption: <string> (optional)
        Define encryption type for smtp: "none", "ssl", "starttls".
        Default is none.
      login: <string> (optional)
        User login for smtp server. If no login is specified, assume no
        login is required.
      password: <string> (optional)
        Password associated to the login.

    Example:
      bounce: {Sender: foo@example.com, To: toto@example.com,
                host: mail.example.com}
    """
    name = "bounce"
    def __init__(self, sender, to, host, port=None, encryption=None, login=None,
                 password=None):
        ImapamiAction.__init__(self, fetch="all")
        self.sender = sender
        self.to = to
        self.smtp = ImapamiSmtpSender(host, port, encryption, login, password)

    def process(self, ami, mail):
        sender = self.evaluate(self.sender, ami, mail.msg)
        to = self.evaluate(self.to, ami, mail.msg)
        return self.smtp.send(ami, sender, to, str(mail.msg))
register(ImapamiActionBounce)

class ImapamiActionSetVar(ImapamiAction):
    """
    Set a variable.

    As it does not read/write the mail, this action does not need
    to fetch the mail by default.
    If a mail header field is used as a variable in the action
    parameters, the user must ensure that the mail headers are fetched,
    for instance by setting the rule paramater 'fetch' to 'headers'.

    Arguments:
      <var name>: <string>
        Each key contains the name of the variable to be set,
        and its value contains the string that should be set. If the
        value is None, the variable is deleted. Multiple fields can be
        specified.

    Example:
      set-var: {var1: foo@example.com, var2: None}
    """
    name = "set-var"
    def __init__(self, **kwargs):
        ImapamiAction.__init__(self)
        self.kwargs = kwargs
    def process(self, ami, mail):
        for k, v in self.kwargs.iteritems():
            k = self.evaluate(k, ami, mail.msg)
            if v is None and ami.variables.get(k) is not None:
                ami.variables.pop(k)
            else:
                v = self.evaluate(v, ami, mail.msg)
                ami.variables[k] = v
        return True
register(ImapamiActionSetVar)

def new(config):
    """
    Create an action object from its yaml config.

    :arg string config:
      The yaml action config.
    :returns:
      The action object.
    """
    logger = logging.getLogger('imapami')
    logger.debug("parsing action %s", config)
    if len(config) != 1:
        raise ValueError("the action config must be a dictionary whose only "
                         "key is the action name")
    action_name = config.keys()[0]
    argslist = []
    argsdict = {}
    if isinstance(config[action_name], list):
        argslist = config[action_name]
    elif isinstance(config[action_name], dict):
        argsdict = config[action_name]
    elif config[action_name] is not None:
        argslist = [config[action_name]]
    action_class = _all_actions.get(action_name)
    if action_class is None:
        raise ValueError("Invalid action name '%s'" % action_name)
    logger.debug("new action %s(%s, %s)", action_name, argslist, argsdict)
    return action_class(*argslist, **argsdict)

def get():
    """
    Return a dictionary containing all the registered actions.
    """
    return _all_actions