Как правильно написать калькулятор на питоне с помощью eval()

В комментариях к статьям по синтаксическому анализу я иногда вижу такие:

на питоне калькулятор пишется проще простого — print(eval(input()))

Ну, вобщем‑то — да, но если, например, вы прикрутите такой калькулятор к своему сайту, то любой желающий вместо 2+2*2 может написать exec("import os; os.removedirs('/')"), предварительно изучив все ваши секретные файлы подобным же образом. Такая перспектива не может радовать, но и отказываться от eval() тоже не стоит.

— А что делать‑то? — спросите вы. Ответ простой: валидировать входящие данные, как вы это всегда делаете. Только не те, которые передаются функции eval() — для этого вам действительно пришлось бы написать свой синтаксический анализатор, а внутренние данные, которыми оперирует реализация eval().

— А! Я знаю — скажете вы — используем compile(), валидируем код и тогда уже вызываем eval(). Что-ж, возможно и так. Но я не знаю простого способа валидировать байт‑код. Нет. Нам надо валидировать результат питоновского синтаксического анализатора — абстрактное синтаксическое дерево или AST. Это гораздо проще, хоть и звучит пугающе.

Напишем свой eval() так:

def my_eval(expression):
    tree = ast.parse(expression, mode='eval')
    code = compile(tree, filename='', mode='eval')
    return eval(code)

Нам, естественнно, понадобится модуль ast:

>>> import ast

Проверяем - работает ли?

>>> print(my_eval(input()))
2+2*2
6

Итак, у нас есть tree. Вам, конечно же, любопытно, что там внутри, и действительно, надо же знать, что мы собираемся валидировать. Но просто так посмотреть не получится — это объект, и print() или pprint()выдадут всего лишь <ast.Expression object at 0x7f6de7eb3c70> Я пользовался такой функцией:

def dump(node):
    def _format(node, indent):
        if isinstance(node, ast.AST):
            print('%sAST %s' % (' ' * indent, node.__class__.__name__))
            for a, b in ast.iter_fields(node):
                print('%s%s' % (' ' * indent, a))
                _format(b, indent + 4)
        elif isinstance(node, list):
            print('%sLIST %s' % (' ' * indent, node.__class__.__name__))
            for x in node:
                _format(x, indent)
        else:
            print('%s%s' % (' ' * indent, repr(node)))
    _format(node, 0)

Можете вставить dump(tree) в my_eval(), поиграться и посмотреть, какое оно — это дерево. Только не надо копипастить мой пример с removedirs().

Кому играться лениво — вот пара примеров:

>>> print(my_eval(input()))
2+2*2
AST Expression
body
    AST BinOp
    left
        AST Constant
        value
            2
        kind
            None
    op
        AST Add
    right
        AST BinOp
        left
            AST Constant
            value
                2
            kind
                None
        op
            AST Mult
        right
            AST Constant
            value
                2
            kind
                None
6

>>> print(my_eval(input()))
print(globals())
AST Expression
body
    AST Call
    func
        AST Name
        id
            'print'
        ctx
            AST Load
    args
        LIST list
        AST Call
        func
            AST Name
            id
                'globals'
            ctx
                AST Load
        args
            LIST list
        keywords
            LIST list
    keywords
        LIST list
{'__name__': '__main__', '__doc__': None, '__package__': None, '__loader__': <class '_frozen_importlib.BuiltinImporter'>, '__spec__': None, '__annotations__': {}, '__builtins__': <module 'builtins' (built-in)>, 'my_eval': <function my_eval at 0x7f6de9225480>, 'ast': <module 'ast' from '/usr/lib/python3.10/ast.py'>, 'dump': <function dump at 0x7f6de7efb490>}
None

Видите, чем отличаются безопасные выражения от небезопасных? Безопасные содержат только Constant и BinOp, а небезопасные — всякие Call и Name. Полный список операций и других узлов дерева можно посмотреть в документации к модулю ast. Это поможет определиться, что запретить, а что разрешить. Для простого калькулятора, я считаю, достаточно разрешить основные операции BinOp и UnaryOp. Плюс Constant.

Пора писать валидатор:

_allowed_nodes = (
    # базовые узлы:
    ast.BinOp, ast.UnaryOp, ast.Constant,

    # основные BinOps:
    ast.Add, ast.Sub, ast.Mult, ast.Div, ast.FloorDiv, ast.Mod, ast.Pow,

    # основные UnaryOps:
    ast.UAdd, ast.USub
)

def validate_ast(tree):

    # валидируем корень дерева
    if not isinstance(tree, ast.Expression):
        raise Exception('Неправильное выражение')

    # валидируем узлы
    def validate_children(node):
        for child in ast.iter_child_nodes(node):
            if not isinstance(child, _allowed_nodes):
                raise Exception('Неправильное выражение')
            validate_children(child)

    validate_children(tree)

Вот так всё просто. Смело вызывайте эту функцию из my_eval() перед compile():

def my_eval(expression):
    tree = ast.parse(expression, mode='eval')
    validate_ast(tree)
    code = compile(tree, filename='', mode='eval')
    return eval(code)

и будет вам счастье. Но, имейте ввиду, документация к функции compile() сообщает:

Warning

It is possible to crash the Python interpreter with a sufficiently large/complex string when compiling to an AST object due to stack depth limitations in Python’s AST compiler.

Так что ограничивайте длину строки, передаваемой my_eval(). И тогда счастье будет настоящим.

— Постойте! — скажете вы — Мне не нужен простой бухгалтерский калькулятор! Где хоть какой-нибудь минимальный набор математических функций? Не отправлять же пользователя в пешее путешествие искать косинус фи?

М-да, пожалуй, соглашусь с вами. Давайте разрешим кое-что из модуля math. Но это будут вызовы функций, а значит — небезопасные области тьмы. Нам придётся добавить ast.Call, ast.Name и ast.Load в _allowed_nodes, и это открывает путь к exec(). Чтобы избежать неприятностей, надо ограничить контекст выполнения, в котором важно наличие пустого __builtins__:

def my_eval(expression):
    tree = ast.parse(expression, mode='eval')
    validate_ast(tree)
    code = compile(tree, filename='', mode='eval')
    context = {'__builtins__': {}}
    return eval(code, context)

Плюс, надо добавить поддержку узлов-списков в валидатор:

def validate_ast(tree):

    # валидируем корень дерева
    if not isinstance(tree, ast.Expression):
        raise Exception('Неправильное выражение')

    # валидируем узлы
    def validate_children(node):
        for child in ast.iter_child_nodes(node):
            if isinstance(child, list):
                for grandchild in child:
                    validate_children(grandchild)
            else:
                if not isinstance(child, _allowed_nodes):
                    raise Exception('Неправильное выражение')
                validate_children(child)

    validate_children(tree)

Теперь, если мы введём

exec('print(42)')

то получим

NameError: name 'exec' is not defined

Ура! Ничего лучшего я, конечно, придумать не могу, но надеюсь что те, кто умнее меня, напишут в комментариях как это взломать и, соответственно, защититься.

Итак, функции у нас разрешены, надо добавить нужные в контекст выполнения. Я ограничусь синусом и косинусом, остальное вы и сами сможете. По своему вкусу.

import math

def my_eval(expression):
    tree = ast.parse(expression, mode='eval')
    validate_ast(tree)
    code = compile(tree, filename='', mode='eval')
    context = {
        '__builtins__': {},
        'sin': math.sin,
        'cos': math.cos
    }
    return eval(code, context)

Может показаться, что достаточно создать контекст всего один раз, вне функции, как мы это сделали с _allowed_nodes в валидаторе, но! — _allowed_nodes у нас всё-таки immutable, а для контекста есть риск модификации, случайной или преднамеренной. Известные аналогичные грабли — использование dict и list в качестве значений по умолчанию для аргументов. Поэтому лучше создавать контекст каждый раз, внутри функции.

Ну, вот собственно, и всё. Пользуйтесь и наслаждайтесь. Для тех, кому лень собирать копипастой всё в кучу, ловите полный исходник:

import ast
import math


def my_eval(expression):
    tree = ast.parse(expression, mode='eval')
    #dump(tree)
    validate_ast(tree)
    code = compile(tree, filename='', mode='eval')
    context = {
        '__builtins__': {},
        'sin': math.sin,
        'cos': math.cos
    }
    return eval(code, context)


_allowed_nodes = (
    # базовые узлы:
    ast.BinOp, ast.UnaryOp, ast.Constant, ast.Call, ast.Name, ast.Load,

    # основные BinOps:
    ast.Add, ast.Sub, ast.Mult, ast.Div, ast.FloorDiv, ast.Mod, ast.Pow,

    # основные UnaryOps:
    ast.UAdd, ast.USub
)

def validate_ast(tree):

    # валидируем корень дерева
    if not isinstance(tree, ast.Expression):
        raise Exception('Неправильное выражение')

    # валидируем узлы
    def validate_children(node):
        for child in ast.iter_child_nodes(node):
            if isinstance(child, list):
                for grandchild in child:
                    validate_children(grandchild)
            else:
                if not isinstance(child, _allowed_nodes):
                    raise Exception('Неправильное выражение')
                validate_children(child)

    validate_children(tree)


def dump(node):
    def _format(node, indent):
        if isinstance(node, ast.AST):
            print('%sAST %s' % (' ' * indent, node.__class__.__name__))
            for a, b in ast.iter_fields(node):
                print('%s%s' % (' ' * indent, a))
                _format(b, indent + 4)
        elif isinstance(node, list):
            print('%sLIST %s' % (' ' * indent, node.__class__.__name__))
            for x in node:
                _format(x, indent)
        else:
            print('%s%s' % (' ' * indent, repr(node)))
    _format(node, 0)


print(my_eval(input()))