#!/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 logging
import re

_all_conditions = {}

def register(cond_class):
    """
    Register a condition class
    """
    if not issubclass(cond_class, ImapamiCond):
        raise ValueError('"%s" condition is not a subclass of ImapamiCond' %
                         cond_class)
    if cond_class.name in _all_conditions:
        raise ValueError('"%s" condition is already registered' %
                         cond_class.name)
    _all_conditions[cond_class.name] = cond_class

class ImapamiCond(object):
    """
    This is the parent class for conditions.

    A condition is a test performed on the mails located in a
    specific mailbox directory. For instance, a condition can check
    the sender, the subject, the date of a mail, or the content of
    a variable.
    """
    def __init__(self, fetch=None, criteria=None):
        """
        Generic condition 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 set() criteria:
          IMAP criteria for IMAP search command. Ex: set(["UNSEEN"])
        """
        self.fetch = fetch or "no"
        self.criteria = criteria or set()

    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 check(self, ami, mail):
        """
        Check the condition.

        :arg Imapami ami:
          The Imapami object
        :arg ImapamiMail mail:
          The mail data
        :returns:
          True if the condition matches, else False.
        """
        return True

    def get_criteria(self):
        """
        Get the list of IMAP criteria that are passed to the server inside
        the IMAP search command.
        The caller will evaluate the variables, so it's not to be done in the
        ImapamiCond class.
        """
        return " ".join([str(c) for c in self.criteria])

class ImapamiCondNot(ImapamiCond):
    """
    Invert the result of a condition.

    Arguments:
      <condition>

    Example:
      not: {from: foo@example.com}
    """
    name = "not"
    def __init__(self, cond):
        cond = new(cond)
        criteria = set(['NOT (%s)' % (cond.get_criteria())])
        ImapamiCond.__init__(self, fetch=cond.fetch, criteria=criteria)
        self.cond = cond

    def check(self, ami, mail):
        return not self.cond.check(ami, mail)
register(ImapamiCondNot)

class ImapamiCondUnseen(ImapamiCond):
    """
    Match if a mail is not marked as seen.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      None

    Example:
      unseen: {}
    """
    name = "unseen"
    def __init__(self):
        ImapamiCond.__init__(self, criteria=set(["UNSEEN"]))
register(ImapamiCondUnseen)

class ImapamiCondFrom(ImapamiCond):
    """
    Match if a 'From' field contains the specified substring.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      substr: <string>
        The substring that should be included in the 'From' header
        to match.

    Example:
      from: {substr: foo@example.com}
    """
    name = "from"
    def __init__(self, substr):
        ImapamiCond.__init__(self, criteria=set(['FROM "%s"' % substr]))
register(ImapamiCondFrom)

class ImapamiCondSubject(ImapamiCond):
    """
    Match if a 'Subject' field contains the specified substring.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      substr: <string>
        The substring that should be included in the 'Subject' header
        to match.

    Example:
      subject: {substr: foo@example.com}
    """
    name = "subject"
    def __init__(self, substr):
        ImapamiCond.__init__(self, criteria=set(['SUBJECT "%s"' % substr]))
register(ImapamiCondSubject)

class ImapamiCondTo(ImapamiCond):
    """
    Match if a 'To' field contains the specified substring.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      substr: <string>
        The substring that should be included in the 'To' header
        to match.

    Example:
      to: {substr: foo@example.com}
    """
    name = "to"
    def __init__(self, substr):
        ImapamiCond.__init__(self, criteria=set(['TO "%s"' % substr]))
register(ImapamiCondTo)

class ImapamiCondCc(ImapamiCond):
    """
    Match if a 'Cc' field contains the specified substring.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      substr: <string>
        The substring that should be included in the 'Cc' header
        to match.

    Example:
      cc: {substr: foo@example.com}
    """
    name = "cc"
    def __init__(self, substr):
        ImapamiCond.__init__(self, criteria=set(['CC "%s"' % substr]))
register(ImapamiCondCc)

class ImapamiCondBcc(ImapamiCond):
    """
    Match if a 'Bcc' field contains the specified substring.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      substr: <string>
        The substring that should be included in the 'Bcc' header
        to match.

    Example:
      bcc: {substr: foo@example.com}
    """
    name = "bcc"
    def __init__(self, substr):
        ImapamiCond.__init__(self, criteria=set(['BCC "%s"' % substr]))
register(ImapamiCondBcc)

class ImapamiCondBody(ImapamiCond):
    """
    Match if the body of the message contains the specified substring.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      substr: <string>
        The substring that should be included in the 'Body' header
        to match.

    Example:
      body: {substr: foobar}
    """
    name = "body"
    def __init__(self, substr):
        ImapamiCond.__init__(self, criteria=set(['BODY "%s"' % substr]))
register(ImapamiCondBody)

class ImapamiCondSince(ImapamiCond):
    """
    Match if the message has been sent since the specified date.

    Match for messages whose internal date (disregarding time and timezone)
    is within or later than the specified date.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      date: <string>
        The reference date. The format is day-month-year, with:
        day = 2 digits, month = 3 letter string (Jan, Feb, Mar, Apr,
        May, Jun, Jul, Aug, Sep, Oct, Nov, Dec), year = 4 digits.

    Example:
      since: {date: 28-Oct-2015}
    """
    name = "since"
    def __init__(self, date):
        ImapamiCond.__init__(self, criteria=set(['SINCE "%s"' % date]))
register(ImapamiCondSince)

class ImapamiCondSentBefore(ImapamiCond):
    """
    Match if 'Date' header is earlier than the specified date (disregarding
    time and timezone).

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      sent-before: <string>
        The reference date. The format is day-month-year, with:
        day = 2 digits, month = 3 letter string (Jan, Feb, Mar, Apr,
        May, Jun, Jul, Aug, Sep, Oct, Nov, Dec), year = 4 digits.

    Example:
      sent-before: {date: 28-Oct-2015}
    """
    name = "sent-before"
    def __init__(self, date):
        ImapamiCond.__init__(self, criteria=set(['SENTBEFORE "%s"' % date]))
register(ImapamiCondSentBefore)

class ImapamiCondSentOn(ImapamiCond):
    """
    Match if 'Date' header is within the specified date (disregarding
    time and timezone).

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      sent-on: <string>
        The reference date. The format is day-month-year, with:
        day = 2 digits, month = 3 letter string (Jan, Feb, Mar, Apr,
        May, Jun, Jul, Aug, Sep, Oct, Nov, Dec), year = 4 digits.

    Example:
      sent-on: {date: 28-Oct-2015}
    """
    name = "sent-on"
    def __init__(self, date):
        ImapamiCond.__init__(self, criteria=set(['SENTON "%s"' % date]))
register(ImapamiCondSentOn)

class ImapamiCondSentSince(ImapamiCond):
    """
    Match if 'Date' header is within or later than the specified date
    (disregarding time and timezone).

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      sent-since: <string>
        The reference date. The format is day-month-year, with:
        day = 2 digits, month = 3 letter string (Jan, Feb, Mar, Apr,
        May, Jun, Jul, Aug, Sep, Oct, Nov, Dec), year = 4 digits.

    Example:
      sent-since: {date: 28-Oct-2015}
    """
    name = "sent-since"
    def __init__(self, date):
        ImapamiCond.__init__(self, criteria=set(['SENTSINCE "%s"' % date]))
register(ImapamiCondSentSince)

class ImapamiCondRecent(ImapamiCond):
    """
    Match if the mail is marked as recent

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      None

    Example:
      recent: {}
    """
    name = "recent"
    def __init__(self, substr):
        ImapamiCond.__init__(self, criteria=set(['RECENT']))
register(ImapamiCondRecent)

class ImapamiCondAnswered(ImapamiCond):
    """
    Match if the mail is not marked as answered.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      None

    Example:
      answered: {}
    """
    name = "answered"
    def __init__(self, substr):
        ImapamiCond.__init__(self, criteria=set(['ANSWERED']))
register(ImapamiCondAnswered)

class ImapamiCondUnanswered(ImapamiCond):
    """
    Match if the mail is not marked as answered.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      None

    Example:
      unanswered: {}
    """
    name = "unanswered"
    def __init__(self, substr):
        ImapamiCond.__init__(self, criteria=set(['UNANSWERED']))
register(ImapamiCondUnanswered)

class ImapamiCondFlagged(ImapamiCond):
    """
    Match if the message is flagged.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      None

    Example:
      flagged: {}
    """
    name = "flagged"
    def __init__(self, substr):
        ImapamiCond.__init__(self, criteria=set(['FLAGGED']))
register(ImapamiCondFlagged)

class ImapamiCondUnflagged(ImapamiCond):
    """
    Match if the message is not flagged.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      None

    Example:
      unflagged: {}
    """
    name = "unflagged"
    def __init__(self, substr):
        ImapamiCond.__init__(self, criteria=set(['UNFLAGGED']))
register(ImapamiCondUnflagged)

class ImapamiCondDraft(ImapamiCond):
    """
    Match if the message is marked as draft.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      None

    Example:
      draft: {}
    """
    name = "draft"
    def __init__(self, substr):
        ImapamiCond.__init__(self, criteria=set(['DRAFT']))
register(ImapamiCondDraft)

class ImapamiCondUndraft(ImapamiCond):
    """
    Match if the message is not marked as draft.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      None

    Example:
      undraft: {}
    """
    name = "undraft"
    def __init__(self, substr):
        ImapamiCond.__init__(self, criteria=set(['UNDRAFT']))
register(ImapamiCondUndraft)

class ImapamiCondKeyword(ImapamiCond):
    """
    Match if the messages has the specified keyword flag set.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      key: <string>
        The keyword string.

    Example:
      keyword: {key: important}
    """
    name = "keyword"
    def __init__(self, key):
        ImapamiCond.__init__(self, criteria=set(['KEYWORD "%s"' % key]))
register(ImapamiCondKeyword)

class ImapamiCondUnkeyword(ImapamiCond):
    """
    Match if the messages does not have the specified keyword flag set.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      key: <string>
        The keyword string.

    Example:
      unkeyword: {key: important}
    """
    name = "unkeyword"
    def __init__(self, key):
        ImapamiCond.__init__(self, criteria=set(['UNKEYWORD "%s"' % key]))
register(ImapamiCondUnkeyword)

class ImapamiCondLarger(ImapamiCond):
    """
    Match if the message is larger than the specified size.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      size: <integer>
        The size of the message in bytes.

    Example:
      larger: {size: 1024}
    """
    name = "larger"
    def __init__(self, size):
        ImapamiCond.__init__(self, criteria=set(['LARGER "%s"' % size]))
register(ImapamiCondLarger)

class ImapamiCondSmaller(ImapamiCond):
    """
    Match if the message is smaller than the specified size.

    This condition does not require to fetch the mail as it can
    be filtered by the IMAP server.

    Arguments:
      size: <integer>
        The size of the message in bytes.

    Example:
      smaller: {size: 1024}
    """
    name = "smaller"
    def __init__(self, size):
        ImapamiCond.__init__(self, criteria=set(['SMALLER "%s"' % size]))
register(ImapamiCondSmaller)

class ImapamiCondRegex(ImapamiCond):
    """
    Match if the regular expression is found in the specified field.

    This condition requires at least to fetch the mail headers, or the
    full message if the regexp is researched in the body of the mail.

    Arguments:
      field: <string>
        The name of the header field (ex: 'From', 'To', ...) where the
        regular expression is researched or 'part1'/'all' if the regular
        expression is researched in the body of the mail ('all' includes
        attachments).
      pattern: <string>
        The regular expression, as supported by the python re module.

    Example:
      regexp: {field: Subject, pattern: '\\[my_list\\]'}
    """
    name = "regexp"
    def __init__(self, field, pattern):
        if field == "part1":
            fetch = "part1"
        elif field == "all":
            fetch = "all"
        else:
            fetch = "headers"
        ImapamiCond.__init__(self, fetch=fetch)
        self.field = field
        self.pattern = pattern
    def check(self, ami, mail):
        field = self.evaluate(self.field, ami, mail.msg)
        pattern = self.evaluate(self.pattern, ami,
                                mail.msg)
        if field == "part1":
            m = re.search(pattern, mail.body_part1)
        elif field == "all":
            m = re.search(pattern, mail.body_all)
        else:
            data = mail.msg.get(field)
            if data is None:
                return False
            m = re.search(pattern, data)
        if m:
            return True
        return False

register(ImapamiCondRegex)

class ImapamiCondAnd(ImapamiCond):
    """
    Match if all conditions of a list match (AND).

    This command is used internally when the list of conditions of a
    rule is parsed. It can also be used by a user to group conditions:
    as soon as a condition does not match, this meta condition returns
    False (lazy evaluation).

    Arguments:
      List of conditions

    Example:
      and:
      - regexp: {field: From, pattern: foo}
      - regexp: {field: Subject, pattern: bar}
    """
    name = "and"
    def __init__(self, *cond_list):
        cond_list = [new(c) for c in cond_list]
        criteria = set().union(*[c.criteria for c in cond_list])
        fetch = imapami.utils.highest_fetch_level(
            [c.fetch for c in cond_list])
        ImapamiCond.__init__(self, fetch=fetch, criteria=criteria)
        self.cond_list = cond_list

    def check(self, ami, mail):
        for c in self.cond_list:
            if c.check(ami, mail) == False:
                return False
        return True
register(ImapamiCondAnd)

class ImapamiCondOr(ImapamiCond):
    """
    Match if at least one condition of a list matches (OR).

    Try to match the conditions of the list: as soon as a condition matches,
    this meta condition returns True (lazy evaluation).

    Arguments:
      List of conditions

    Example:
      or:
      - regexp: {field: From, pattern: foo}
      - regexp: {field: Subject, pattern: bar}
    """
    name = "or"
    def __init__(self, *cond_list):
        cond_list = [new(c) for c in cond_list]
        criteria = ''
        for c in cond_list:
            crit = c.get_criteria()
            if crit == '':
                continue
            if criteria == '':
                criteria = crit
            else:
                criteria = 'OR (%s) (%s)' % (criteria, crit)
        if criteria != '':
            criteria = set([criteria])
        fetch = imapami.utils.highest_fetch_level(
            [c.fetch for c in cond_list])
        ImapamiCond.__init__(self, fetch=fetch, criteria=criteria)
        self.cond_list = cond_list

    def check(self, ami, mail):
        for c in self.cond_list:
            if c.check(ami, mail) == True:
                return True
        return False
register(ImapamiCondOr)

class ImapamiCondEq(ImapamiCond):
    """
    Match if strings are equal.

    This condition does not fetch any part of the mail by default.
    If a mail header field is used as a variable in the condition
    parameters, the user must ensure that the mail headers are fetched,
    for instance by setting the rule paramater 'fetch' to 'headers'.

    If more than 2 elements are given in the list, all of them must
    be equal.

    Arguments:
      List of strings

    Example:
      eq: ['foo@example.com', '{From}', '{foo}']
    """
    name = "eq"
    def __init__(self, *str_list):
        ImapamiCond.__init__(self)
        if not isinstance(str_list, list) and not isinstance(str_list, tuple):
            raise ValueError("arguments of 'eq' should be a list/tuple")
        if len(str_list) < 2:
            raise ValueError("'eq' argument list is too short")
        self.str_list = list(str_list)

    def check(self, ami, mail):
        first = self.evaluate(self.str_list[0], ami,
                              mail.msg)
        for s in self.str_list[1:]:
            if self.evaluate(s, ami, mail.msg) != first:
                return False
        return True
register(ImapamiCondEq)

def new(config):
    """
    Create a condition object from its yaml config.

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

def get():
    """
    Return a dictionary containing all the registered conditions.
    """
    return _all_conditions
