# Proof-of-concept Jinja2 conditional expression evaluator.


# Based on code from Ansible, under the following license agreement:
#
# (c) 2012-2014, Michael DeHaan <michael dehaan gmail com>
#
# This file is part of Ansible
#
# Ansible 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 3 of the License, or
# (at your option) any later version.
#
# Ansible is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.


import ast
import re

import jinja2


class Jinja2ExpressionEvaluator:
    def __init__(self, variables):
        self.variables = variables

        self.environment = jinja2.Environment()

        # Remove the default global functions such as range(), dict(), lipsum()
        # etc. -- these don't make sense for conditional expressions.
        self.environment.globals = []

    def template(self, data):
        '''Expand the given string using the Jinja2 templating engine'''
        template = self.environment.from_string(data)
        new_context = template.new_context(self.variables, shared=True)
        result = template.root_render_func(new_context)
        return jinja2.utils.concat(result)

    DEFINED_REGEX = re.compile(r'([\w_]+)\s+(not\s+is|is|is\s+not)\s+(defined|undefined)')
    def extract_defined_undefined(self, conditional):
        '''Extract subexpressions like '$var is defined' from the expression.

        We need to treat these specially because undefined variables trigger an
        exception from the rendering function.

        '''
        results = []

        cond = conditional
        m = self.DEFINED_REGEX.search(cond)
        while m:
            results.append(m.groups())
            cond = cond[m.end():]
            m = self.DEFINED_REGEX.search(cond)

        return result

    def evaluate_conditional(self, conditional):
        try:
            presented = "{%% if %s %%} True {%% else %%} False {%% endif %%}" % conditional
            val = self.template(presented).strip()
            if val == "True":
                return True
            elif val == "False":
                return False
            else:
                raise RuntimeError("unable to evaluate conditional: %s" % original)
        except jinja2.exceptions.UndefinedError as e:
            # The templating failed, meaning most likely a variable was undefined. If we happened
            # to be looking for an undefined variable, return True, otherwise fail
            try:
                # first we extract the variable name from the error message
                var_name = re.compile(r"'(hostvars\[.+\]|[\w_]+)' is undefined").search(str(e)).groups()[0]
                # next we extract all defined/undefined tests from the conditional string
                def_undef = self.extract_defined_undefined(conditional)
                # then we loop through these, comparing the error variable name against
                # each def/undef test we found above. If there is a match, we determine
                # whether the logic/state mean the variable should exist or not and return
                # the corresponding True/False
                for (du_var, logic, state) in def_undef:
                    # when we compare the var names, normalize quotes because something
                    # like hostvars['foo'] may be tested against hostvars["foo"]
                    if var_name.replace("'", '"') == du_var.replace("'", '"'):
                        # the should exist is a xor test between a negation in the logic portion
                        # against the state (defined or undefined)
                        should_exist = ('not' in logic) != (state == 'defined')
                        if should_exist:
                            return False
                        else:
                            return True
                # as nothing above matched the failed var name, re-raise here to
                # trigger the UndefinedVariable exception again below
                raise
            except Exception:
                raise jinja2.exceptions.UndefinedError("error while evaluating conditional (%s): %s" % (original, e))


def evaluate_condition(text, variables):
    engine = Jinja2ExpressionEvaluator(variables=variables)
    return engine.evaluate_conditional(text)


variables = {
    'xyzzy': 3,
    'plugh': 10,
    'plover': 'string',
    'y2': True,
    'features': ['first', 'second', ['nested', 'thing']],
}

tests = [
    # Simple boolean logic
    '9 + 5 == 14',
    # Variable evaluations
    'plugh == 9',
    'xyzzy > 2 and plugh > 9',
    'plover[1] == "s"',
    'plover == "String"',
    # Type coercion
    'y2',
    'plover',
    'y2 == True',
    'y2 == "True"',
    '5 == "5"',
    '0',
    # The 'defined' function takes a special codepath that triggers when
    # jinja2 raises an UndefinedError exception.
    'y2 is defined',
    'y3 is defined',
    '(y3 is defined or plover == "false") or True',
    # Lists
    '"second" in features',
    '"third" in features',
    'features[0] == "second"',
    '"foo" in ["foo", "bar", "baz"]',
]

for condition in tests:
    result = evaluate_condition(condition, variables)
    print("%s: %s" % (condition, result))