BNIL Expression Visitor

BNILExprVisitor.py

class BNILExprVisitor(object):
    '''A class to faciliate visiting BNIL instructions.

    The following example outputs all addition expressions that are assigned
    to an MLIL variable.

    >>> visit = BNILExprVisitor()
    >>> @visit.add(MediumLevelILOperation.MLIL_SET_VAR)
    ... def visit_set_var(expr)
    ...     visit(expr.src)
    >>> @visit.add(MediumLevelILOperation.MLIL_ADD)
    ... def visit_add(expr):
    ...     print repr(expr)
    >>> for bb in current_function.medium_level_il:
    ...     for i in bb:
    ...         visit(i)
    <il: rax + 8>
    <il: rax + 8>
    <il: rax + 8>
    <il: var_2c8 + rax.eax>
    <il: var_2c0 + 1>
    <il: var_2c4 + 1>
    '''
    def __init__(self):
        self._visitors = dict()
        self._default_visitor = None

    def add(self, operation):
        '''Decorator for adding a visitor function for a BNIL operation.
        '''
        def add_decorator(visitor):
            self._visitors[operation] = visitor
            return visitor
        return add_decorator

    def add_default(self, visitor):
        '''Decorator to add a default visitor.
        '''
        self._default_visitor = visitor
        return visitor

    def __call__(self, expr, *args, **kwargs):
        try:
            visitor = self._visitors[expr.operation]
            return visitor(expr, *args, **kwargs)
        except KeyError:
            if self._default_visitor:
                return self._default_visitor(expr, *args, **kwargs)
            else:
                return