File size: 1,989 Bytes
9c1c0ef
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
from __future__ import annotations

import ast

BLOCKED_NODES = (
    ast.Import,
    ast.ImportFrom,
    ast.Global,
    ast.Nonlocal,
    ast.With,
    ast.AsyncWith,
    ast.Try,
    ast.Raise,
    ast.ClassDef,
    ast.FunctionDef,
    ast.AsyncFunctionDef,
)
BLOCKED_CALLS = {
    "eval",
    "exec",
    "compile",
    "open",
    "input",
    "__import__",
    "breakpoint",
    "getattr",
    "setattr",
    "delattr",
}


class UnsafeCodeError(ValueError):
    pass


def validate_generated_expression(expression: str, max_characters: int = 2_000) -> ast.Expression:
    """Validate a calculation expression before it is sent to an isolated worker.



    DataPilot's primary workflow does not execute LLM-generated Python. This validator

    exists for an optional restricted calculation worker and accepts expressions only.

    """
    if len(expression) > max_characters:
        raise UnsafeCodeError("Expression exceeds the configured size limit.")
    try:
        tree = ast.parse(expression, mode="eval")
    except SyntaxError as exc:
        raise UnsafeCodeError("Expression is not valid Python.") from exc
    for node in ast.walk(tree):
        if isinstance(node, BLOCKED_NODES):
            raise UnsafeCodeError(f"Blocked syntax: {type(node).__name__}.")
        if isinstance(node, ast.Attribute) and node.attr.startswith("__"):
            raise UnsafeCodeError("Dunder attribute access is blocked.")
        if isinstance(node, ast.Name) and node.id.startswith("__"):
            raise UnsafeCodeError("Dunder names are blocked.")
        if isinstance(node, ast.Call):
            if isinstance(node.func, ast.Name) and node.func.id in BLOCKED_CALLS:
                raise UnsafeCodeError(f"Blocked call: {node.func.id}.")
            if isinstance(node.func, ast.Attribute) and node.func.attr in BLOCKED_CALLS:
                raise UnsafeCodeError(f"Blocked call: {node.func.attr}.")
    return tree