210 lines
9.1 KiB
Python
210 lines
9.1 KiB
Python
"""Small data-only expression interpreter. Never executes Python bytecode."""
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import math
|
|
import operator
|
|
import re
|
|
import string
|
|
import time
|
|
from functools import lru_cache
|
|
from typing import Any
|
|
|
|
MAX_TEXT = 16384
|
|
MAX_ITEMS = 4096
|
|
MAX_NODES = 384
|
|
MAX_DEPTH = 24
|
|
|
|
|
|
class ExpressionError(ValueError):
|
|
pass
|
|
|
|
|
|
def bounded(value: Any, depth: int = 0, budget: list[int] | None = None) -> Any:
|
|
if budget is None:
|
|
budget = [MAX_TEXT * 4]
|
|
if depth > MAX_DEPTH:
|
|
raise ExpressionError("Value nesting limit")
|
|
budget[0] -= 1
|
|
if type(value) is str:
|
|
budget[0] -= len(value)
|
|
if len(value) > MAX_TEXT:
|
|
raise ExpressionError("Text limit")
|
|
elif type(value) is int:
|
|
if value.bit_length() > 256:
|
|
raise ExpressionError("Integer limit")
|
|
elif type(value) is float:
|
|
if not math.isfinite(value):
|
|
raise ExpressionError("Finite numbers required")
|
|
elif value is None or type(value) is bool:
|
|
pass
|
|
elif type(value) in (list, tuple, dict):
|
|
if len(value) > MAX_ITEMS:
|
|
raise ExpressionError("Collection limit")
|
|
for item in value:
|
|
bounded(item, depth + 1, budget)
|
|
if type(value) is dict:
|
|
bounded(value[item], depth + 1, budget)
|
|
else:
|
|
raise ExpressionError("Only plain data values are allowed")
|
|
if budget[0] < 0:
|
|
raise ExpressionError("Value size limit")
|
|
return value
|
|
|
|
|
|
_BINARY = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul,
|
|
ast.Div: operator.truediv, ast.FloorDiv: operator.floordiv, ast.Mod: operator.mod,
|
|
ast.Pow: operator.pow}
|
|
_COMPARE = {ast.Eq: operator.eq, ast.NotEq: operator.ne, ast.Lt: operator.lt,
|
|
ast.LtE: operator.le, ast.Gt: operator.gt, ast.GtE: operator.ge,
|
|
ast.In: lambda a, b: a in b, ast.NotIn: lambda a, b: a not in b,
|
|
ast.Is: operator.is_, ast.IsNot: operator.is_not}
|
|
_ALLOWED = (ast.Expression, ast.Constant, ast.Name, ast.Load, ast.BinOp, ast.UnaryOp,
|
|
ast.BoolOp, ast.Compare, ast.IfExp, ast.Call, ast.keyword, ast.List, ast.Tuple,
|
|
ast.Dict, ast.Subscript, ast.Slice, ast.And, ast.Or, ast.Not, ast.UAdd, ast.USub,
|
|
*_BINARY.keys(), *_COMPARE.keys())
|
|
|
|
|
|
@lru_cache(maxsize=256)
|
|
def parse_expression(expression: str) -> ast.Expression:
|
|
if len(expression) > 4096:
|
|
raise ExpressionError("Expression length limit")
|
|
tree = ast.parse(expression, mode="eval")
|
|
nodes = list(ast.walk(tree))
|
|
if len(nodes) > MAX_NODES:
|
|
raise ExpressionError("Expression complexity limit")
|
|
for node in nodes:
|
|
if not isinstance(node, _ALLOWED):
|
|
raise ExpressionError("Unsupported expression operation")
|
|
if isinstance(node, ast.Name) and node.id.startswith("__"):
|
|
raise ExpressionError("Private names are forbidden")
|
|
if isinstance(node, ast.Call) and (not isinstance(node.func, ast.Name)
|
|
or len(node.args) + len(node.keywords) > 32):
|
|
raise ExpressionError("Only named helper calls are allowed")
|
|
if isinstance(node, ast.keyword) and node.arg is None:
|
|
raise ExpressionError("Argument unpacking is forbidden")
|
|
if isinstance(node, ast.Dict) and any(key is None for key in node.keys):
|
|
raise ExpressionError("Dictionary unpacking is forbidden")
|
|
return tree
|
|
|
|
|
|
def bounded_template(template: Any, row: dict[str, Any]) -> str:
|
|
"""Literal field lookup; no attribute traversal or unbounded format widths."""
|
|
text = str(template)
|
|
bounded(text)
|
|
result = []
|
|
for literal, name, spec, conversion in string.Formatter().parse(text):
|
|
result.append(literal)
|
|
if name is None:
|
|
continue
|
|
if conversion not in (None, "s", "r") or not re.fullmatch(r"(?:[<^>])?\d{0,3}(?:\.\d{1,2})?[dfs]?", spec or ""):
|
|
raise ExpressionError("Unsupported template format")
|
|
value = bounded(row.get(name, ""))
|
|
if conversion == "r":
|
|
value = repr(value)
|
|
elif conversion == "s":
|
|
value = str(value)
|
|
result.append(format(value, spec or ""))
|
|
return bounded("".join(result))
|
|
|
|
|
|
def bounded_replace(text: str, old: str, new: str = "") -> str:
|
|
for value in (text, old, new):
|
|
bounded(value)
|
|
size = len(text) + text.count(old) * (len(new) - len(old))
|
|
if size > MAX_TEXT:
|
|
raise ExpressionError("Replacement size limit")
|
|
return text.replace(old, new)
|
|
|
|
|
|
def evaluate_expression(expression: str, helpers: dict[str, Any], values: dict[str, Any]) -> Any:
|
|
tree = parse_expression(expression)
|
|
remaining = MAX_NODES
|
|
deadline = time.monotonic() + 0.05
|
|
|
|
def visit(node: ast.AST, depth: int = 0) -> Any:
|
|
nonlocal remaining
|
|
remaining -= 1
|
|
if remaining < 0 or depth > MAX_DEPTH or time.monotonic() > deadline:
|
|
raise ExpressionError("Expression evaluation limit")
|
|
child = lambda item: visit(item, depth + 1)
|
|
if isinstance(node, ast.Expression):
|
|
result = child(node.body)
|
|
elif isinstance(node, ast.Constant):
|
|
result = node.value
|
|
elif isinstance(node, ast.Name):
|
|
# Functions are never returned as values or passed as arguments.
|
|
result = values.get(node.id, "")
|
|
elif isinstance(node, (ast.List, ast.Tuple)):
|
|
result = [child(item) for item in node.elts]
|
|
if isinstance(node, ast.Tuple):
|
|
result = tuple(result)
|
|
elif isinstance(node, ast.Dict):
|
|
result = {child(key): child(value) for key, value in zip(node.keys, node.values)}
|
|
elif isinstance(node, ast.UnaryOp):
|
|
operand = child(node.operand)
|
|
result = not operand if isinstance(node.op, ast.Not) else (+operand if isinstance(node.op, ast.UAdd) else -operand)
|
|
elif isinstance(node, ast.BoolOp):
|
|
result = child(node.values[0])
|
|
for item in node.values[1:]:
|
|
if (isinstance(node.op, ast.And) and not result) or (isinstance(node.op, ast.Or) and result):
|
|
break
|
|
result = child(item)
|
|
elif isinstance(node, ast.IfExp):
|
|
result = child(node.body if child(node.test) else node.orelse)
|
|
elif isinstance(node, ast.Compare):
|
|
left = child(node.left)
|
|
result = True
|
|
for op, right_node in zip(node.ops, node.comparators):
|
|
right = child(right_node)
|
|
if not _COMPARE[type(op)](left, right):
|
|
result = False
|
|
break
|
|
left = right
|
|
elif isinstance(node, ast.BinOp):
|
|
left, right = child(node.left), child(node.right)
|
|
if isinstance(node.op, ast.Mult):
|
|
sequence, count = (left, right) if type(right) is int else (right, left)
|
|
if type(sequence) in (str, list, tuple) and type(count) is int:
|
|
if len(sequence) * max(0, count) > (MAX_TEXT if type(sequence) is str else MAX_ITEMS):
|
|
raise ExpressionError("Repeat size limit")
|
|
if isinstance(node.op, ast.Mod) and type(left) is str:
|
|
raise ExpressionError("Use concat or tpl for text formatting")
|
|
if isinstance(node.op, ast.Pow) and (type(right) not in (int, float) or abs(right) > 16 or abs(left) > 1e12):
|
|
raise ExpressionError("Power limit")
|
|
result = _BINARY[type(node.op)](left, right)
|
|
elif isinstance(node, ast.Subscript):
|
|
container = child(node.value)
|
|
if type(container) not in (dict, list, tuple, str):
|
|
raise ExpressionError("Only data can be indexed")
|
|
if isinstance(node.slice, ast.Slice):
|
|
parts = [child(part) if part is not None else None for part in
|
|
(node.slice.lower, node.slice.upper, node.slice.step)]
|
|
index = slice(*parts)
|
|
else:
|
|
index = child(node.slice)
|
|
result = container[index]
|
|
elif isinstance(node, ast.Call):
|
|
name = node.func.id
|
|
fn = helpers.get(name)
|
|
if name.startswith("_") or not callable(fn):
|
|
raise ExpressionError("Unknown helper")
|
|
args = [child(arg) for arg in node.args]
|
|
kwargs = {item.arg: child(item.value) for item in node.keywords}
|
|
# Prevent expansion before constructing the result in text helpers.
|
|
if name in {"swap", "thousands", "format_thousands"}:
|
|
for value in [*args[1:], *kwargs.values()]:
|
|
if type(value) is str and len(value) > 64:
|
|
raise ExpressionError("Separator size limit")
|
|
if name in {"format_time", "time_format", "fmt_time", "format_distance", "distance",
|
|
"time_diff", "format_time_diff", "number", "num"}:
|
|
for value in [*args[1:], *kwargs.values()]:
|
|
if type(value) is int and abs(value) > 100:
|
|
raise ExpressionError("Format precision limit")
|
|
result = fn(*args, **kwargs)
|
|
else:
|
|
raise ExpressionError("Unsupported expression")
|
|
return bounded(result)
|
|
|
|
return visit(tree)
|