# expr.py - expression handling functions

#

# Copyright (C) 2011, 2012, 2013 Arthur de Jong

#

# This library is free software; you can redistribute it and/or

# modify it under the terms of the GNU Lesser General Public

# License as published by the Free Software Foundation; either

# version 2.1 of the License, or (at your option) any later version.

#

# This library 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

# Lesser General Public License for more details.

#

# You should have received a copy of the GNU Lesser General Public

# License along with this library; if not, write to the Free Software

# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA

# 02110-1301 USA

"""Module for handling expressions used for LDAP searches.

>>> expr = Expression('foo=$foo')

>>> expr.value(dict(foo='XX'))

'foo=XX'

>>> expr = Expression('foo=${foo:-$bar}')

>>> expr.value(dict(foo='', bar='YY'))

'foo=YY'

>>> expr.value(dict(bar=['YY', 'ZZ']))

'foo=YY'

>>> Expression(r'${passwd#{crypt\}}').value(dict(passwd='{crypt}HASH'))

'HASH'

>>> Expression('${var#trim}').value(dict(var='notrimme'))

'notrimme'

>>> Expression('${var#?trim}').value(dict(var='xtrimme'))

'me'

>>> Expression('${var#*trim}').value(dict(var='xxxtrimme'))

'me'

>>> Expression('${var%.txt}').value(dict(var='foo.txt'))

'foo'

>>> Expression('${x#$y}').value(dict(x='a/b', y='a'))

'/b'

>>> Expression('${var#t*is}').value(dict(var='this is a test'))

' is a test'

>>> Expression('${var##t*is}').value(dict(var='this is a test'))

' a test'

>>> Expression('${var%t*st}').value(dict(var='this is a test'))

'this is a '

>>> Expression('${var%%t*st}').value(dict(var='this is a test'))

''

>>> Expression('${test1:0:6}').value(dict(test1='foobar'))

'foobar'

>>> Expression('${test1:0:10}').value(dict(test1='foobar'))

'foobar'

>>> Expression('${test1:0:3}').value(dict(test1='foobar'))

'foo'

>>> Expression('${test1:3:0}').value(dict(test1='foobar'))

''

>>> Expression('${test1:3:6}').value(dict(test1='foobar'))

'bar'

>>> Expression('${test1:7:0}').value(dict(test1='foobar'))

''

>>> Expression('${test1:7:3}').value(dict(test1='foobar'))

''

"""

import fnmatch

import re

# exported names

__all__ = ('Expression', )

# TODO: do more expression validity checking

class MyIter(object):

    """Custom iterator-like class with a back() method."""

    def __init__(self, value):

        self.value = value

        self.pos = 0

    def next(self):

        self.pos += 1

        try:

            return self.value[self.pos - 1]

        except IndexError:

            return None

    def startswith(self, value):

        return self.value[self.pos].startswith(value)

    def back(self):

        self.pos -= 1

    def __iter__(self):

        return self

    def get_name(self):

        """Read a variable name from the value iterator."""

        name = ''

        for c in self:

            if not c or not c.isalnum():

                self.back()

                return name

            name += c

        return name

class DollarExpression(object):

    """Class for handling a variable $xxx ${xxx}, ${xxx:-yyy} or ${xxx:+yyy}

    expression."""

    def __init__(self, value):

        """Parse the expression as the start of a $-expression."""

        self.op = None

        self.expr = None

        c = value.next()

        if c == '{':

            self.name = value.get_name()

            c = value.next()

            if c == '}':

                return

            elif c == ':':

                if value.startswith('-') or value.startswith('+'):

                    # ${attr:-word} or ${attr:+word}

                    self.op = c + value.next()

                else:

                    # ${attr:offset:length}

                    self.op = c

            elif c in ('#', '%'):

                c2 = value.next()

                if c2 in ('#', '%'):

                    c += c2

                else:

                    value.back()

                self.op = c

            else:

                raise ValueError('Expecting operator')

            self.expr = Expression(value, endat='}')

        elif c == '(':

            self.name = None

            self.op = value.get_name()

            c = value.next()

            if c != '(':

                raise ValueError("Expecting '('")

            self.expr = Expression(value, endat=')')

            c = value.next()

            if c != ')':

                raise ValueError("Expecting ')'")

        else:

            value.back()

            self.name = value.get_name()

    def value(self, variables):

        """Expand the expression using the variables specified."""

        # lookup the value

        value = variables.get(self.name, '')

        if value in (None, [], ()):

            value = ''

        elif isinstance(value, (list, tuple)):

            value = value[0]

        # TODO: try to return multiple values, one for each value of the list

        if self.op == ':-':

            return value if value else self.expr.value(variables)

        elif self.op == ':+':

            return self.expr.value(variables) if value else ''

        elif self.op == ':':

            offset, length = self.expr.value(variables).split(':')

            offset, length = int(offset), int(length)

            return value[offset:offset + length]

        elif self.op in ('#', '##', '%', '%%'):

            match = fnmatch.translate(self.expr.value(variables))

            if self.op == '#':

                match = match.replace('*', '*?').replace(r'\Z', r'(?P<replace>.*)\Z')

            elif self.op == '##':

                match = match.replace(r'\Z', r'(?P<replace>.*?)\Z')

            elif self.op == '%':

                match = r'(?P<replace>.*)' + match.replace('*', '*?')

            elif self.op == '%%':

                match = r'(?P<replace>.*?)' + match

            match = re.match(match, value)

            return match.group('replace') if match else value

        elif self.op == 'lower':

            return self.expr.value(variables).lower()

        elif self.op == 'upper':

            return self.expr.value(variables).upper()

        return value

    def variables(self, results):

        """Add the variables used in the expression to results."""

        if self.name:

            results.add(self.name)

        if self.expr:

            self.expr.variables(results)

class Expression(object):

    """Class for parsing and expanding an expression."""

    def __init__(self, value, endat=None):

        """Parse the expression as a string."""

        if not isinstance(value, MyIter):

            self.expression = value

            value = MyIter(value)

        expr = []

        literal = ''

        c = value.next()

        while c != endat:

            if c == '$':

                if literal:

                    expr.append(literal)

                expr.append(DollarExpression(value))

                literal = ''

            elif c == '\\':

                literal += value.next()

            else:

                literal += c

            c = value.next()

        if literal:

            expr.append(literal)

        self.expr = expr

    def value(self, variables):

        """Expand the expression using the variables specified."""

        res = ''

        for x in self.expr:

            if hasattr(x, 'value'):

                res += x.value(variables)

            else:

                res += x

        return res

    def variables(self, results=None):

        """Return the variables defined in the expression."""

        if not results:

            results = set()

        for x in self.expr:

            if hasattr(x, 'variables'):

                x.variables(results)

        return results

    def __str__(self):

        return self.expression

    def __repr__(self):

        return repr(str(self))