26894310350

Committed 03 Jun 2026 03:15PM UTC coverage: 98.993% (+0.02%) from 98.972%

Build # 26894310350

Build Type

Pull #317

github

Committed by

web-flow

Commit Message

Merge 22ced9021 into a2891c0d1

Pull Request Pull Request #317: Type Annotations for RestrictedPython

Coverage Stats

212 of 219 branches covered (96.8%)

120 of 120 new or added lines in 8 files covered. (100.0%)

2 existing lines in 1 file now uncovered.

2556 of 2582 relevant lines covered (98.99%)

0.99 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

99.57

/src/RestrictedPython/transformer.py

##############################################################################
#
# Copyright (c) 2002 Zope Foundation and Contributors.
#
# This software is subject to the provisions of the Zope Public License,
# Version 2.1 (ZPL).  A copy of the ZPL should accompany this distribution.
# THIS SOFTWARE IS PROVIDED "AS IS" AND ANY AND ALL EXPRESS OR IMPLIED
# WARRANTIES ARE DISCLAIMED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF TITLE, MERCHANTABILITY, AGAINST INFRINGEMENT, AND FITNESS
# FOR A PARTICULAR PURPOSE
#
##############################################################################
"""
transformer module:

uses Python standard library ast module and its containing classes to transform
the parsed python code to create a modified AST for a byte code generation.
"""


import ast
import collections
import contextlib
import sys
import textwrap
import typing


# For AugAssign the operator must be converted to a string.
IOPERATOR_TO_STR = {
    ast.Add: '+=',
    ast.Sub: '-=',
    ast.Mult: '*=',
    ast.Div: '/=',
    ast.Mod: '%=',
    ast.Pow: '**=',
    ast.LShift: '<<=',
    ast.RShift: '>>=',
    ast.BitOr: '|=',
    ast.BitXor: '^=',
    ast.BitAnd: '&=',
    ast.FloorDiv: '//=',
    ast.MatMult: '@=',
}

# For creation allowed magic method names. See also
# https://docs.python.org/3/reference/datamodel.html#special-method-names
ALLOWED_FUNC_NAMES = frozenset([
    '__init__',
    '__contains__',
    '__lt__',
    '__le__',
    '__eq__',
    '__ne__',
    '__gt__',
    '__ge__',
])


FORBIDDEN_FUNC_NAMES = frozenset([
    'print',
    'printed',
    'builtins',
    'breakpoint',
])

# Attributes documented in the `inspect` module, but defined on the listed
# objects. See also https://docs.python.org/3/library/inspect.html
INSPECT_ATTRIBUTES = frozenset([
    # on traceback objects:
    "tb_frame",
    # "tb_lasti",  # int
    # "tb_lineno",  # int
    "tb_next",
    # on frame objects:
    "f_back",
    "f_builtins",
    "f_code",
    "f_generator",
    "f_globals",
    # "f_lasti",  # int
    # "f_lineno",  # int
    "f_locals",
    "f_trace",
    # on code objects:
    # "co_argcount",  # int
    "co_code",
    # "co_cellvars",  # tuple of str
    # "co_consts",   # tuple of str
    # "co_filename",  # str
    # "co_firstlineno",  # int
    # "co_flags",  # int
    # "co_lnotab",  # mapping between ints and indices
    # "co_freevars",  # tuple of strings
    # "co_posonlyargcount",  # int
    # "co_kwonlyargcount",  # int
    # "co_name",  # str
    # "co_qualname",  # str
    # "co_names",  # str
    # "co_nlocals",  # int
    # "co_stacksize",  # int
    # "co_varnames",  # tuple of str
    # on generator objects:
    "gi_frame",
    # "gi_running",  # bool
    # "gi_suspended",  # bool
    "gi_code",
    "gi_yieldfrom",
    # on coroutine objects:
    "cr_await",
    "cr_frame",
    # "cr_running",  # bool
    "cr_code",
    "cr_origin",
])

_T_visit_return: typing.TypeAlias = ast.AST | typing.Iterable[ast.AST] | None
_T_pos_ast: typing.TypeAlias = (
    ast.stmt | ast.expr | ast.excepthandler | ast.arg | ast.keyword | ast.alias
    | ast.pattern)
if sys.version_info >= (3, 12):
    _T_pos_ast: typing.TypeAlias = _T_pos_ast | ast.type_param
_T = typing.TypeVar('_T', bound=ast.AST)

# When new ast nodes are generated they have no 'lineno', 'end_lineno',
# 'col_offset' and 'end_col_offset'. This function copies these fields from the
# incoming node:


def copy_locations(new_node: _T_pos_ast, old_node: _T_pos_ast) -> None:
    assert 'lineno' in new_node._attributes
    new_node.lineno = old_node.lineno

    assert 'end_lineno' in new_node._attributes
    new_node.end_lineno = old_node.end_lineno

    assert 'col_offset' in new_node._attributes
    new_node.col_offset = old_node.col_offset

    assert 'end_col_offset' in new_node._attributes
    new_node.end_col_offset = old_node.end_col_offset

    ast.fix_missing_locations(new_node)


class PrintInfo:
    def __init__(self) -> None:
        self.print_used = False
        self.printed_used = False

    @contextlib.contextmanager
    def new_print_scope(self) -> collections.abc.Iterator[None]:
        old_print_used = self.print_used
        old_printed_used = self.printed_used

        self.print_used = False
        self.printed_used = False

        try:
            yield
        finally:
            self.print_used = old_print_used
            self.printed_used = old_printed_used


class RestrictingNodeTransformer(ast.NodeTransformer):
    errors: list[str]
    warnings: list[str]
    used_names: dict[str, bool]

    def __init__(self,
                 errors: list[str] | None = None,
                 warnings: list[str] | None = None,
                 used_names: dict[str, bool] | None = None):
        super().__init__()
        self.errors = [] if errors is None else errors
        self.warnings = [] if warnings is None else warnings

        # All the variables used by the incoming source.
        # Internal names/variables, like the ones from 'gen_tmp_name', don't
        # have to be added.
        # 'used_names' is for example needed by 'RestrictionCapableEval' to
        # know wich names it has to supply when calling the final code.
        self.used_names = {} if used_names is None else used_names

        # Global counter to construct temporary variable names.
        self._tmp_idx = 0

        self.print_info = PrintInfo()

    def gen_tmp_name(self) -> str:
        # 'check_name' ensures that no variable is prefixed with '_'.
        # => Its safe to use '_tmp..' as a temporary variable.
        name = '_tmp%i' % self._tmp_idx
        self._tmp_idx += 1
        return name

    def error(self, node: ast.AST, info: str) -> None:
        """Record a security error discovered during transformation."""
        lineno = getattr(node, 'lineno', None)
        self.errors.append(
            f'Line {lineno}: {info}')

    def warn(self, node: ast.AST, info: str) -> None:
        """Record a security warning discovered during transformation."""
        lineno = getattr(node, 'lineno', None)
        self.warnings.append(
            f'Line {lineno}: {info}')

    def guard_iter(self, node: ast.For | ast.comprehension) -> _T_visit_return:
        """
        Converts:
            for x in expr
        to
            for x in _getiter_(expr)

        Also used for
        * list comprehensions
        * dict comprehensions
        * set comprehensions
        * generator expresions
        """
        node = self.node_contents_visit(node)

        if isinstance(node.target, ast.Tuple):
            spec = self.gen_unpack_spec(node.target)
            new_iter = ast.Call(
                func=ast.Name('_iter_unpack_sequence_', ast.Load()),
                args=[node.iter, spec, ast.Name('_getiter_', ast.Load())],
                keywords=[])
        else:
            new_iter = ast.Call(
                func=ast.Name("_getiter_", ast.Load()),
                args=[node.iter],
                keywords=[])

        copy_locations(new_iter, node.iter)
        node.iter = new_iter
        return node

    def is_starred(self, ob: ast.AST) -> typing.TypeGuard[ast.Starred]:
        # TODO: Change Type Annotation to typing.TypeIs[ast.Starred] when
        #               Support for Python 3.12 is dropped.
        return isinstance(ob, ast.Starred)

    def gen_unpack_spec(self, tpl: ast.Tuple) -> ast.Dict:
        """Generate a specification for 'guarded_unpack_sequence'.

        This spec is used to protect sequence unpacking.
        The primary goal of this spec is to tell which elements in a sequence
        are sequences again. These 'child' sequences have to be protected
        again.

        For example there is a sequence like this:
            (a, (b, c), (d, (e, f))) = g

        On a higher level the spec says:
            - There is a sequence of len 3
            - The element at index 1 is a sequence again with len 2
            - The element at index 2 is a sequence again with len 2
              - The element at index 1 in this subsequence is a sequence again
                with len 2

        With this spec 'guarded_unpack_sequence' does something like this for
        protection (len checks are omitted):

            t = list(_getiter_(g))
            t[1] = list(_getiter_(t[1]))
            t[2] = list(_getiter_(t[2]))
            t[2][1] = list(_getiter_(t[2][1]))
            return t

        The 'real' spec for the case above is then:
            spec = {
                'min_len': 3,
                'childs': (
                    (1, {'min_len': 2, 'childs': ()}),
                    (2, {
                            'min_len': 2,
                            'childs': (
                                (1, {'min_len': 2, 'childs': ()})
                            )
                        }
                    )
                )
            }

        So finally the assignment above is converted into:
            (a, (b, c), (d, (e, f))) = guarded_unpack_sequence(g, spec)
        """
        spec = ast.Dict(keys=[], values=[])

        spec.keys.append(ast.Constant('childs'))
        val0 = ast.Tuple([], ast.Load())
        spec.values.append(val0)

        # starred elements in a sequence do not contribute into the min_len.
        # For example a, b, *c = g
        # g must have at least 2 elements, not 3. 'c' is empyt if g has only 2.
        min_len = len([ob for ob in tpl.elts if not self.is_starred(ob)])
        offset = 0

        for idx, val in enumerate(tpl.elts):
            # After a starred element specify the child index from the back.
            # Since it is unknown how many elements from the sequence are
            # consumed by the starred element.
            # For example a, *b, (c, d) = g
            # Then (c, d) has the index '-1'
            if self.is_starred(val):
                offset = min_len + 1

            elif isinstance(val, ast.Tuple):
                el = ast.Tuple([], ast.Load())
                el.elts.append(ast.Constant(idx - offset))
                el.elts.append(self.gen_unpack_spec(val))
                val0.elts.append(el)

        spec.keys.append(ast.Constant('min_len'))
        spec.values.append(ast.Constant(min_len))

        return spec

    def protect_unpack_sequence(
            self,
            target: ast.Tuple,
            value: ast.expr) -> ast.Call:
        spec = self.gen_unpack_spec(target)
        return ast.Call(
            func=ast.Name('_unpack_sequence_', ast.Load()),
            args=[value, spec, ast.Name('_getiter_', ast.Load())],
            keywords=[])

    def gen_unpack_wrapper(self,
                           node: ast.stmt,
                           target: ast.Tuple) -> tuple[ast.Name, ast.Try]:
        """Helper function to protect tuple unpacks.

        node: used to copy the locations for the new nodes.
        target: is the tuple which must be protected.

        It returns a tuple with two element.

        Element 1: Is a temporary name node which must be used to
                   replace the target.
                   The context (store, param) is defined
                   by the 'ctx' parameter..

        Element 2: Is a try .. finally where the body performs the
                   protected tuple unpack of the temporary variable
                   into the original target.
        """

        # Generate a tmp name to replace the tuple with.
        tmp_name = self.gen_tmp_name()

        # Generates an expressions which protects the unpack.
        # converter looks like 'wrapper(tmp_name)'.
        # 'wrapper' takes care to protect sequence unpacking with _getiter_.
        converter = self.protect_unpack_sequence(
            target,
            ast.Name(tmp_name, ast.Load()))

        # Assign the expression to the original names.
        # Cleanup the temporary variable.
        # Generates:
        # try:
        #     # converter is 'wrapper(tmp_name)'
        #     arg = converter
        # finally:
        #     del tmp_arg
        try_body: list[ast.stmt] = [ast.Assign(
            targets=[target], value=converter)]
        finalbody: list[ast.stmt] = [self.gen_del_stmt(tmp_name)]
        cleanup = ast.Try(
            body=try_body, finalbody=finalbody, handlers=[], orelse=[])

        # This node is used to catch the tuple in a tmp variable.
        tmp_target = ast.Name(tmp_name, ast.Store())

        copy_locations(tmp_target, node)
        copy_locations(cleanup, node)

        return (tmp_target, cleanup)

    def gen_none_node(self) -> ast.Constant:
        return ast.Constant(None)

    def gen_del_stmt(self, name_to_del: str) -> ast.Delete:
        return ast.Delete(targets=[ast.Name(name_to_del, ast.Del())])

    def check_name(
            self,
            node: _T_pos_ast,
            name: str | None,
            allow_magic_methods: bool = False) -> None:
        """Check names if they are allowed.

        If ``allow_magic_methods is True`` names in `ALLOWED_FUNC_NAMES`
        are additionally allowed although their names start with `_`.

        """
        if name is None:
            return

        if (name.startswith('_')
                and name != '_'
                and not (allow_magic_methods
                         and name in ALLOWED_FUNC_NAMES
                         and node.col_offset != 0)):
            self.error(
                node,
                '"{name}" is an invalid variable name because it '
                'starts with "_"'.format(name=name))
        elif name.endswith('__roles__'):
            self.error(node, '"%s" is an invalid variable name because '
                       'it ends with "__roles__".' % name)
        elif name in FORBIDDEN_FUNC_NAMES:
            self.error(node, f'"{name}" is a reserved name.')

    def check_function_argument_names(
            self,
            node: ast.FunctionDef | ast.AsyncFunctionDef | ast.Lambda) -> None:
        for arg in node.args.args:
            self.check_name(node, arg.arg)

        if node.args.vararg:
            self.check_name(node, node.args.vararg.arg)

        if node.args.kwarg:
            self.check_name(node, node.args.kwarg.arg)

        for arg in node.args.kwonlyargs:
            self.check_name(node, arg.arg)

    def check_import_names(self, node: ast.ImportFrom | ast.Import) -> ast.AST:
        """Check the names being imported.

        This is a protection against rebinding dunder names like
        _getitem_, _write_ via imports.

        => 'from _a import x' is ok, because '_a' is not added to the scope.
        """
        for name in node.names:
            if '*' in name.name:
                self.error(node, '"*" imports are not allowed.')
            self.check_name(node, name.name)
            if name.asname:
                self.check_name(node, name.asname)

        return self.node_contents_visit(node)

    def inject_print_collector(
            self,
            node: ast.Module | ast.FunctionDef,
            position: int = 0) -> None:
        print_used = self.print_info.print_used
        printed_used = self.print_info.printed_used

        if print_used or printed_used:
            # Add '_print = _print_(_getattr_)' add the top of a
            # function/module.
            _print = ast.Assign(
                targets=[ast.Name('_print', ast.Store())],
                value=ast.Call(
                    func=ast.Name("_print_", ast.Load()),
                    args=[ast.Name("_getattr_", ast.Load())],
                    keywords=[]))

            if isinstance(node, ast.Module):
                _print.lineno = position
                _print.col_offset = position
                _print.end_lineno = position
                _print.end_col_offset = position
                ast.fix_missing_locations(_print)
            else:
                copy_locations(_print, node)

            node.body.insert(position, _print)

            if not printed_used:
                self.warn(node, "Prints, but never reads 'printed' variable.")

            elif not print_used:
                self.warn(node, "Doesn't print, but reads 'printed' variable.")

    # Special Functions for an ast.NodeTransformer

    def generic_visit(self,  # type: ignore[override]
                      node: ast.AST) -> _T_visit_return:
        """Reject ast nodes which do not have a corresponding `visit_` method.

        This is needed to prevent new ast nodes from new Python versions to be
        trusted before any security review.

        To access `generic_visit` on the super class use `node_contents_visit`.
        """
        self.warn(
            node,
            '{0.__class__.__name__}'
            ' statement is not known to RestrictedPython'.format(node)
        )
        self.not_allowed(node)

    def not_allowed(self, node: ast.AST) -> None:
        self.error(
            node,
            f'{node.__class__.__name__} statements are not allowed.')

    def node_contents_visit(self, node: _T) -> _T:
        """Visit the contents of a node."""
        return super().generic_visit(node)  # type: ignore[return-value]

    # ast for Literals

    def visit_Constant(self, node: ast.Constant) -> _T_visit_return:
        """Allow constant literals.

        Constant replaces Num, Str, Bytes, NameConstant and Ellipsis in
        Python 3.8+.
        :see: https://docs.python.org/dev/whatsnew/3.8.html#deprecated
        """
        return self.node_contents_visit(node)

    def visit_Interactive(self, node: ast.Interactive) -> _T_visit_return:
        """Allow single mode without restrictions."""
        return self.node_contents_visit(node)

    def visit_List(self, node: ast.List) -> _T_visit_return:
        """Allow list literals without restrictions."""
        return self.node_contents_visit(node)

    def visit_Tuple(self, node: ast.Tuple) -> _T_visit_return:
        """Allow tuple literals without restrictions."""
        return self.node_contents_visit(node)

    def visit_Set(self, node: ast.Set) -> _T_visit_return:
        """Allow set literals without restrictions."""
        return self.node_contents_visit(node)

    def visit_Dict(self, node: ast.Dict) -> _T_visit_return:
        """Allow dict literals without restrictions."""
        return self.node_contents_visit(node)

    def visit_FormattedValue(
            self,
            node: ast.FormattedValue) -> _T_visit_return:
        """Allow f-strings without restrictions."""
        return self.node_contents_visit(node)

    def visit_TemplateStr(self, node: ast.AST) -> _T_visit_return:
        """Template strings are allowed by default.

        As Template strings are a very basic template mechanism, that needs
        additional rendering logic to be useful, they are not blocked by
        default.
        Those rendering logic would be affected by RestrictedPython as well.

        TODO: Change Type Annotation to ast.TemplateStr when
              Support for Python 3.13 is dropped.
        """
        return self.node_contents_visit(node)

    def visit_Interpolation(self, node: ast.AST) -> _T_visit_return:
        """Interpolations are allowed by default.

        As Interpolations are part of Template Strings, they are needed
        to be reached in the context of RestrictedPython as Template Strings
        are allowed. As a user has to provide additional rendering logic
        to make use of Template Strings, the security implications of
        Interpolations are limited in the context of RestrictedPython.

        TODO: Change Type Annotation to ast.Interpolation when
              Support for Python 3.13 is dropped.
        """
        return self.node_contents_visit(node)

    def visit_JoinedStr(self, node: ast.JoinedStr) -> _T_visit_return:
        """Allow joined string without restrictions."""
        return self.node_contents_visit(node)

    # ast for Variables

    def visit_Name(self, node: ast.Name) -> _T_visit_return:
        """Prevents access to protected names.

        Converts use of the name 'printed' to this expression: '_print()'
        """

        node = self.node_contents_visit(node)

        if isinstance(node.ctx, ast.Load):
            new_node: _T_pos_ast
            if node.id == 'printed':
                self.print_info.printed_used = True
                new_node = ast.Call(
                    func=ast.Name("_print", ast.Load()),
                    args=[],
                    keywords=[])

                copy_locations(new_node, node)
                return new_node

            elif node.id == 'print':
                self.print_info.print_used = True
                new_node = ast.Attribute(
                    value=ast.Name('_print', ast.Load()),
                    attr="_call_print",
                    ctx=ast.Load())

                copy_locations(new_node, node)
                return new_node

            self.used_names[node.id] = True

        self.check_name(node, node.id)
        return node

    def visit_Load(self, node: ast.Load) -> _T_visit_return:
        """

        """
        return self.node_contents_visit(node)

    def visit_Store(self, node: ast.Store) -> _T_visit_return:
        """

        """
        return self.node_contents_visit(node)

    def visit_Del(self, node: ast.Del) -> _T_visit_return:
        """

        """
        return self.node_contents_visit(node)

    def visit_Starred(self, node: ast.Starred) -> _T_visit_return:
        """

        """
        return self.node_contents_visit(node)

    # Expressions

    def visit_Expression(self, node: ast.Expression) -> _T_visit_return:
        """Allow Expression statements without restrictions.

        They are in the AST when using the `eval` compile mode.
        """
        return self.node_contents_visit(node)

    def visit_Expr(self, node: ast.Expr) -> _T_visit_return:
        """Allow Expr statements (any expression) without restrictions."""
        return self.node_contents_visit(node)

    def visit_UnaryOp(self, node: ast.UnaryOp) -> _T_visit_return:
        """
        UnaryOp (Unary Operations) is the overall element for:
        * Not --> which should be allowed
        * UAdd --> Positive notation of variables (e.g. +var)
        * USub --> Negative notation of variables (e.g. -var)
        """
        return self.node_contents_visit(node)

    def visit_UAdd(self, node: ast.UAdd) -> _T_visit_return:
        """Allow positive notation of variables. (e.g. +var)"""
        return self.node_contents_visit(node)

    def visit_USub(self, node: ast.USub) -> _T_visit_return:
        """Allow negative notation of variables. (e.g. -var)"""
        return self.node_contents_visit(node)

    def visit_Not(self, node: ast.Not) -> _T_visit_return:
        """Allow the `not` operator."""
        return self.node_contents_visit(node)

    def visit_Invert(self, node: ast.Invert) -> _T_visit_return:
        """Allow `~` expressions."""
        return self.node_contents_visit(node)

    def visit_BinOp(self, node: ast.BinOp) -> _T_visit_return:
        """Allow binary operations."""
        return self.node_contents_visit(node)

    def visit_Add(self, node: ast.Add) -> _T_visit_return:
        """Allow `+` expressions."""
        return self.node_contents_visit(node)

    def visit_Sub(self, node: ast.Sub) -> _T_visit_return:
        """Allow `-` expressions."""
        return self.node_contents_visit(node)

    def visit_Mult(self, node: ast.Mult) -> _T_visit_return:
        """Allow `*` expressions."""
        return self.node_contents_visit(node)

    def visit_Div(self, node: ast.Div) -> _T_visit_return:
        """Allow `/` expressions."""
        return self.node_contents_visit(node)

    def visit_FloorDiv(self, node: ast.FloorDiv) -> _T_visit_return:
        """Allow `//` expressions."""
        return self.node_contents_visit(node)

    def visit_Mod(self, node: ast.Mod) -> _T_visit_return:
        """Allow `%` expressions."""
        return self.node_contents_visit(node)

    def visit_Pow(self, node: ast.Pow) -> _T_visit_return:
        """Allow `**` expressions."""
        return self.node_contents_visit(node)

    def visit_LShift(self, node: ast.LShift) -> _T_visit_return:
        """Allow `<<` expressions."""
        return self.node_contents_visit(node)

    def visit_RShift(self, node: ast.RShift) -> _T_visit_return:
        """Allow `>>` expressions."""
        return self.node_contents_visit(node)

    def visit_BitOr(self, node: ast.BitOr) -> _T_visit_return:
        """Allow `|` expressions."""
        return self.node_contents_visit(node)

    def visit_BitXor(self, node: ast.BitXor) -> _T_visit_return:
        """Allow `^` expressions."""
        return self.node_contents_visit(node)

    def visit_BitAnd(self, node: ast.BitAnd) -> _T_visit_return:
        """Allow `&` expressions."""
        return self.node_contents_visit(node)

    def visit_MatMult(self, node: ast.MatMult) -> _T_visit_return:
        """Allow multiplication (`@`)."""
        return self.node_contents_visit(node)

    def visit_BoolOp(self, node: ast.BoolOp) -> _T_visit_return:
        """Allow bool operator without restrictions."""
        return self.node_contents_visit(node)

    def visit_And(self, node: ast.And) -> _T_visit_return:
        """Allow bool operator `and` without restrictions."""
        return self.node_contents_visit(node)

    def visit_Or(self, node: ast.Or) -> _T_visit_return:
        """Allow bool operator `or` without restrictions."""
        return self.node_contents_visit(node)

    def visit_Compare(self, node: ast.Compare) -> _T_visit_return:
        """Allow comparison expressions without restrictions."""
        return self.node_contents_visit(node)

    def visit_Eq(self, node: ast.Eq) -> _T_visit_return:
        """Allow == expressions."""
        return self.node_contents_visit(node)

    def visit_NotEq(self, node: ast.NotEq) -> _T_visit_return:
        """Allow != expressions."""
        return self.node_contents_visit(node)

    def visit_Lt(self, node: ast.Lt) -> _T_visit_return:
        """Allow < expressions."""
        return self.node_contents_visit(node)

    def visit_LtE(self, node: ast.LtE) -> _T_visit_return:
        """Allow <= expressions."""
        return self.node_contents_visit(node)

    def visit_Gt(self, node: ast.Gt) -> _T_visit_return:
        """Allow > expressions."""
        return self.node_contents_visit(node)

    def visit_GtE(self, node: ast.GtE) -> _T_visit_return:
        """Allow >= expressions."""
        return self.node_contents_visit(node)

    def visit_Is(self, node: ast.Is) -> _T_visit_return:
        """Allow `is` expressions."""
        return self.node_contents_visit(node)

    def visit_IsNot(self, node: ast.IsNot) -> _T_visit_return:
        """Allow `is not` expressions."""
        return self.node_contents_visit(node)

    def visit_In(self, node: ast.In) -> _T_visit_return:
        """Allow `in` expressions."""
        return self.node_contents_visit(node)

    def visit_NotIn(self, node: ast.NotIn) -> _T_visit_return:
        """Allow `not in` expressions."""
        return self.node_contents_visit(node)

    def visit_Call(self, node: ast.Call) -> _T_visit_return:
        """Checks calls with '*args' and '**kwargs'.

        Note: The following happens only if '*args' or '**kwargs' is used.

        Transfroms 'foo(<all the possible ways of args>)' into
        _apply_(foo, <all the possible ways for args>)

        The thing is that '_apply_' has only '*args', '**kwargs', so it gets
        Python to collapse all the myriad ways to call functions
        into one manageable from.

        From there, '_apply_()' wraps args and kws in guarded accessors,
        then calls the function, returning the value.
        """

        if isinstance(node.func, ast.Name):
            if node.func.id == 'exec':
                self.error(node, 'Exec calls are not allowed.')
            elif node.func.id == 'eval':
                self.error(node, 'Eval calls are not allowed.')

        needs_wrap = False

        for pos_arg in node.args:
            if isinstance(pos_arg, ast.Starred):
                needs_wrap = True

        for keyword_arg in node.keywords:
            if keyword_arg.arg is None:
                needs_wrap = True

        node = self.node_contents_visit(node)

        if not needs_wrap:
            return node

        node.args.insert(0, node.func)
        node.func = ast.Name('_apply_', ast.Load())
        copy_locations(node.func, node.args[0])
        return node

    def visit_keyword(self, node: ast.keyword) -> _T_visit_return:
        """

        """
        return self.node_contents_visit(node)

    def visit_IfExp(self, node: ast.IfExp) -> _T_visit_return:
        """Allow `if` expressions without restrictions."""
        return self.node_contents_visit(node)

    def visit_Attribute(self, node: ast.Attribute) -> _T_visit_return:
        """Checks and mutates attribute access/assignment.

        'a.b' becomes '_getattr_(a, "b")'
        'a.b = c' becomes '_write_(a).b = c'
        'del a.b' becomes 'del _write_(a).b'

        The _write_ function should return a security proxy.
        """
        if node.attr.startswith('_') and node.attr != '_':
            self.error(
                node,
                '"{name}" is an invalid attribute name because it starts '
                'with "_".'.format(name=node.attr))

        if node.attr.endswith('__roles__'):
            self.error(
                node,
                '"{name}" is an invalid attribute name because it ends '
                'with "__roles__".'.format(name=node.attr))

        if node.attr in INSPECT_ATTRIBUTES:
            self.error(
                node,
                f'"{node.attr}" is a restricted name,'
                ' that is forbidden to access in RestrictedPython.',
            )

        if isinstance(node.ctx, ast.Load):
            node = self.node_contents_visit(node)
            new_node = ast.Call(
                func=ast.Name('_getattr_', ast.Load()),
                args=[node.value, ast.Constant(node.attr)],
                keywords=[])

            copy_locations(new_node, node)
            return new_node

        elif isinstance(node.ctx, (ast.Store, ast.Del)):
            node = self.node_contents_visit(node)
            new_value = ast.Call(
                func=ast.Name('_write_', ast.Load()),
                args=[node.value],
                keywords=[])

            copy_locations(new_value, node.value)
            node.value = new_value
            return node

        else:  # pragma: no cover
            # Impossible Case only ctx Load, Store and Del are defined in ast.
            raise NotImplementedError(
                f"Unknown ctx type: {type(node.ctx)}")

    # Subscripting

    def visit_Subscript(self, node: ast.Subscript) -> _T_visit_return:
        """Transforms all kinds of subscripts.

        'foo[bar]' becomes '_getitem_(foo, bar)'
        'foo[:ab]' becomes '_getitem_(foo, slice(None, ab, None))'
        'foo[ab:]' becomes '_getitem_(foo, slice(ab, None, None))'
        'foo[a:b]' becomes '_getitem_(foo, slice(a, b, None))'
        'foo[a:b:c]' becomes '_getitem_(foo, slice(a, b, c))'
        'foo[a, b:c] becomes '_getitem_(foo, (a, slice(b, c, None)))'
        'foo[a] = c' becomes '_write_(foo)[a] = c'
        'del foo[a]' becomes 'del _write_(foo)[a]'

        The _write_ function should return a security proxy.
        """
        node = self.node_contents_visit(node)

        # 'AugStore' and 'AugLoad' are defined in 'Python.asdl' as possible
        # 'expr_context'. However, according to Python/ast.c
        # they are NOT used by the implementation => No need to worry here.
        # Instead ast.c creates 'AugAssign' nodes, which can be visited.

        if isinstance(node.ctx, ast.Load):
            new_node = ast.Call(
                func=ast.Name('_getitem_', ast.Load()),
                args=[node.value, node.slice],
                keywords=[])

            copy_locations(new_node, node)
            return new_node

        elif isinstance(node.ctx, (ast.Del, ast.Store)):
            new_value = ast.Call(
                func=ast.Name('_write_', ast.Load()),
                args=[node.value],
                keywords=[])

            copy_locations(new_value, node)
            node.value = new_value
            return node

        else:  # pragma: no cover
            # Impossible Case only ctx Load, Store and Del are defined in ast.
            raise NotImplementedError(
                f"Unknown ctx type: {type(node.ctx)}")

    def visit_Slice(self, node: ast.Slice) -> _T_visit_return:
        """

        """
        return self.node_contents_visit(node)

    # Comprehensions

    def visit_ListComp(self, node: ast.ListComp) -> _T_visit_return:
        """

        """
        return self.node_contents_visit(node)

    def visit_SetComp(self, node: ast.SetComp) -> _T_visit_return:
        """

        """
        return self.node_contents_visit(node)

    def visit_GeneratorExp(self, node: ast.GeneratorExp) -> _T_visit_return:
        """

        """
        return self.node_contents_visit(node)

    def visit_DictComp(self, node: ast.DictComp) -> _T_visit_return:
        """

        """
        return self.node_contents_visit(node)

    def visit_comprehension(self, node: ast.comprehension) -> _T_visit_return:
        """

        """
        return self.guard_iter(node)

    # Statements

    def visit_Assign(self, node: ast.Assign) -> _T_visit_return:
        """

        """

        node = self.node_contents_visit(node)

        if not any(isinstance(t, ast.Tuple) for t in node.targets):
            return node

        # Handle sequence unpacking.
        # For briefness this example omits cleanup of the temporary variables.
        # Check 'transform_tuple_assign' how its done.
        #
        # - Single target (with nested support)
        # (a, (b, (c, d))) = <exp>
        # is converted to
        # (a, t1) = _getiter_(<exp>)
        # (b, t2) = _getiter_(t1)
        # (c, d) = _getiter_(t2)
        #
        # - Multi targets
        # (a, b) = (c, d) = <exp>
        # is converted to
        # (c, d) = _getiter_(<exp>)
        # (a, b) = _getiter_(<exp>)
        # Why is this valid ? The original bytecode for this multi targets
        # behaves the same way.

        # ast.NodeTransformer works with list results.
        # He injects it at the right place of the node's parent statements.
        new_nodes = []

        # python fills the right most target first.
        for target in reversed(node.targets):
            if isinstance(target, ast.Tuple):
                wrapper = ast.Assign(
                    targets=[target],
                    value=self.protect_unpack_sequence(target, node.value))
                new_nodes.append(wrapper)
            else:
                new_node = ast.Assign(targets=[target], value=node.value)
                new_nodes.append(new_node)

        for new_node in new_nodes:
            copy_locations(new_node, node)

        return new_nodes

    def visit_AugAssign(self, node: ast.AugAssign) -> _T_visit_return:
        """Forbid certain kinds of AugAssign

        According to the language reference (and ast.c) the following nodes
        are are possible:
        Name, Attribute, Subscript

        Note that although augmented assignment of attributes and
        subscripts is disallowed, augmented assignment of names (such
        as 'n += 1') is allowed.
        'n += 1' becomes 'n = _inplacevar_("+=", n, 1)'
        """

        node = self.node_contents_visit(node)

        if isinstance(node.target, ast.Attribute):
            self.error(
                node,
                "Augmented assignment of attributes is not allowed.")
            return node

        elif isinstance(node.target, ast.Subscript):
            self.error(
                node,
                "Augmented assignment of object items "
                "and slices is not allowed.")
            return node

        elif isinstance(node.target, ast.Name):
            new_node = ast.Assign(
                targets=[node.target],
                value=ast.Call(
                    func=ast.Name('_inplacevar_', ast.Load()),
                    args=[
                        ast.Constant(IOPERATOR_TO_STR[type(node.op)]),
                        ast.Name(node.target.id, ast.Load()),
                        node.value
                    ],
                    keywords=[]))

            copy_locations(new_node, node)
            return new_node
        else:  # pragma: no cover
            # Impossible Case - Only Node Types:
            # * Name
            # * Attribute
            # * Subscript
            # defined, those are checked before.
            raise NotImplementedError(
                f"Unknown target type: {type(node.target)}")

    def visit_Raise(self, node: ast.Raise) -> _T_visit_return:
        """Allow `raise` statements without restrictions."""
        return self.node_contents_visit(node)

    def visit_Assert(self, node: ast.Assert) -> _T_visit_return:
        """Allow assert statements without restrictions."""
        return self.node_contents_visit(node)

    def visit_Delete(self, node: ast.Delete) -> _T_visit_return:
        """Allow `del` statements without restrictions."""
        return self.node_contents_visit(node)

    def visit_Pass(self, node: ast.Pass) -> _T_visit_return:
        """Allow `pass` statements without restrictions."""
        return self.node_contents_visit(node)

    # Imports

    def visit_Import(self, node: ast.Import) -> _T_visit_return:
        """Allow `import` statements with restrictions.
        See check_import_names."""
        return self.check_import_names(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> _T_visit_return:
        """Allow `import from` statements with restrictions.
        See check_import_names."""
        return self.check_import_names(node)

    def visit_alias(self, node: ast.alias) -> _T_visit_return:
        """Allow `as` statements in import and import from statements."""
        return self.node_contents_visit(node)

    # Control flow

    def visit_If(self, node: ast.If) -> _T_visit_return:
        """Allow `if` statements without restrictions."""
        return self.node_contents_visit(node)

    def visit_For(self, node: ast.For) -> _T_visit_return:
        """Allow `for` statements with some restrictions."""
        return self.guard_iter(node)

    def visit_While(self, node: ast.While) -> _T_visit_return:
        """Allow `while` statements."""
        return self.node_contents_visit(node)

    def visit_Break(self, node: ast.Break) -> _T_visit_return:
        """Allow `break` statements without restrictions."""
        return self.node_contents_visit(node)

    def visit_Continue(self, node: ast.Continue) -> _T_visit_return:
        """Allow `continue` statements without restrictions."""
        return self.node_contents_visit(node)

    def visit_Try(self, node: ast.Try) -> _T_visit_return:
        """Allow `try` without restrictions."""
        return self.node_contents_visit(node)

    def visit_TryStar(self, node: ast.AST) -> _T_visit_return:
        """Disallow `ExceptionGroup` due to a potential sandbox escape.

        TODO: Type Annotation for node when dropping support
              for Python < 3.11 should be ast.TryStar.
        """
        self.not_allowed(node)

    def visit_ExceptHandler(self, node: ast.ExceptHandler) -> _T_visit_return:
        """Protect exception handlers."""
        node = self.node_contents_visit(node)
        self.check_name(node, node.name)
        return node

    def visit_With(self, node: ast.With) -> _T_visit_return:
        """Protect tuple unpacking on with statements."""
        node = self.node_contents_visit(node)

        for item in reversed(node.items):
            if isinstance(item.optional_vars, ast.Tuple):
                tmp_target, unpack = self.gen_unpack_wrapper(
                    node,
                    item.optional_vars)

                item.optional_vars = tmp_target
                node.body.insert(0, unpack)

        return node

    def visit_withitem(self, node: ast.withitem) -> _T_visit_return:
        """Allow `with` statements (context managers) without restrictions."""
        return self.node_contents_visit(node)

    # Function and class definitions

    def visit_FunctionDef(self, node: ast.FunctionDef) -> _T_visit_return:
        """Allow function definitions (`def`) with some restrictions."""
        self.check_name(node, node.name, allow_magic_methods=True)
        self.check_function_argument_names(node)

        with self.print_info.new_print_scope():
            node = self.node_contents_visit(node)
            self.inject_print_collector(node)
        return node

    def visit_Lambda(self, node: ast.Lambda) -> _T_visit_return:
        """Allow lambda with some restrictions."""
        self.check_function_argument_names(node)
        return self.node_contents_visit(node)

    def visit_arguments(self, node: ast.arguments) -> _T_visit_return:
        """

        """
        return self.node_contents_visit(node)

    def visit_arg(self, node: ast.arg) -> _T_visit_return:
        """

        """
        return self.node_contents_visit(node)

    def visit_Return(self, node: ast.Return) -> _T_visit_return:
        """Allow `return` statements without restrictions."""
        return self.node_contents_visit(node)

    def visit_Yield(self, node: ast.Yield) -> _T_visit_return:
        """Allow `yield`statements without restrictions."""
        return self.node_contents_visit(node)

    def visit_YieldFrom(self, node: ast.YieldFrom) -> _T_visit_return:
        """Allow `yield`statements without restrictions."""
        return self.node_contents_visit(node)

    def visit_Global(self, node: ast.Global) -> _T_visit_return:
        """Allow `global` statements without restrictions."""
        return self.node_contents_visit(node)

    def visit_Nonlocal(self, node: ast.Nonlocal) -> _T_visit_return:
        """Deny `nonlocal` statements."""
        self.not_allowed(node)

    def visit_ClassDef(self, node: ast.ClassDef) -> _T_visit_return:
        """Check the name of a class definition."""
        self.check_name(node, node.name)
        node = self.node_contents_visit(node)
        if any(keyword.arg == 'metaclass' for keyword in node.keywords):
            self.error(
                node, 'The keyword argument "metaclass" is not allowed.')
        CLASS_DEF = textwrap.dedent('''\
            class {0.name}(metaclass=__metaclass__):
                pass
        '''.format(node))
        new_class_node = typing.cast(
            ast.ClassDef, ast.parse(CLASS_DEF).body[0])
        new_class_node.body = node.body
        new_class_node.bases = node.bases
        new_class_node.decorator_list = node.decorator_list
        return new_class_node

    def visit_Module(self, node: ast.Module) -> _T_visit_return:
        """Add the print_collector (only if print is used) at the top."""
        node = self.node_contents_visit(node)

        # Inject the print collector after 'from __future__ import ....'
        position = 0
        for position, child in enumerate(node.body):  # pragma: no branch
            if not isinstance(child, ast.ImportFrom):
                break

            if not child.module == '__future__':
                break

        self.inject_print_collector(node, position)
        return node

    # Async und await

    def visit_AsyncFunctionDef(
            self, node: ast.AsyncFunctionDef) -> _T_visit_return:
        """Deny async functions."""
        self.not_allowed(node)

    def visit_Await(self, node: ast.Await) -> _T_visit_return:
        """Deny async functionality."""
        self.not_allowed(node)

    def visit_AsyncFor(self, node: ast.AsyncFor) -> _T_visit_return:
        """Deny async functionality."""
        self.not_allowed(node)

    def visit_AsyncWith(self, node: ast.AsyncWith) -> _T_visit_return:
        """Deny async functionality."""
        self.not_allowed(node)

    # Assignment expressions (walrus operator ``:=``)
    # New in 3.8
    def visit_NamedExpr(self, node: ast.NamedExpr) -> _T_visit_return:
        """Allow assignment expressions under some circumstances."""
        # while the grammar requires ``node.target`` to be a ``Name``
        # the abstract syntax is more permissive and allows an ``expr``.
        # We support only a ``Name``.
        # This is safe as the expression can only add/modify local
        # variables. While this may hide global variables, an
        # (implicitly performed) name check guarantees (as usual)
        # that no essential global variable is hidden.
        node = self.node_contents_visit(node)  # this checks ``node.target``
        target = node.target
        if not isinstance(target, ast.Name):
            self.error(  # type: ignore[unreachable]
                node,
                "Assignment expressions are only allowed for simple targets")
        return node

zopefoundation / RestrictedPython / 26894310350

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous