obdash/obdcore/formula.py

"""Safe formula evaluator for vehicle-profile PID scaling.

Profiles are community-contributed data, so decode formulas must NOT be able to
execute arbitrary code. Formulas are arithmetic expressions over named
variables -- the de-facto OBD convention used by Torque / FORScan / ScanGauge:

    raw-mode PIDs: variables A, B, C, ... = response data bytes 0, 1, 2, ...
        e.g.  "(A*256+B)*0.57"   "A-40"   "(A>>1)&1"   "A//2"
    derived PIDs:  variables are other PID keys
        e.g.  "MAP - BARO"

Only numeric literals, the named variables, arithmetic/bitwise operators, and a
small whitelist of functions are allowed. No names, attributes, subscripts,
comprehensions, or calls outside the whitelist -- anything else raises
FormulaError at compile time, so a bad/hostile profile fails loudly on load.
"""
import ast
import operator

_BIN = {
    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,
    ast.BitAnd: operator.and_, ast.BitOr: operator.or_, ast.BitXor: operator.xor,
    ast.LShift: operator.lshift, ast.RShift: operator.rshift,
}
_UNARY = {ast.USub: operator.neg, ast.UAdd: operator.pos, ast.Invert: operator.invert}
_FUNCS = {"min": min, "max": max, "abs": abs, "round": round,
          "int": int, "float": float}


class FormulaError(ValueError):
    pass


def _validate(node, allowed):
    if isinstance(node, ast.Expression):
        return _validate(node.body, allowed)
    if isinstance(node, ast.BinOp):
        if type(node.op) not in _BIN:
            raise FormulaError(f"operator not allowed: {type(node.op).__name__}")
        _validate(node.left, allowed)
        _validate(node.right, allowed)
        return
    if isinstance(node, ast.UnaryOp):
        if type(node.op) not in _UNARY:
            raise FormulaError(f"unary op not allowed: {type(node.op).__name__}")
        _validate(node.operand, allowed)
        return
    if isinstance(node, ast.Constant):
        if not isinstance(node.value, (int, float)) or isinstance(node.value, bool):
            raise FormulaError("only numeric constants allowed")
        return
    if isinstance(node, ast.Name):
        if node.id not in allowed:
            raise FormulaError(f"unknown variable {node.id!r} (allowed: {sorted(allowed)})")
        return
    if isinstance(node, ast.Call):
        if not isinstance(node.func, ast.Name) or node.func.id not in _FUNCS:
            raise FormulaError("only min/max/abs/round/int/float calls allowed")
        if node.keywords:
            raise FormulaError("keyword args not allowed")
        for a in node.args:
            _validate(a, allowed)
        return
    raise FormulaError(f"expression not allowed: {type(node).__name__}")


def _eval(node, names):
    if isinstance(node, ast.Expression):
        return _eval(node.body, names)
    if isinstance(node, ast.BinOp):
        return _BIN[type(node.op)](_eval(node.left, names), _eval(node.right, names))
    if isinstance(node, ast.UnaryOp):
        return _UNARY[type(node.op)](_eval(node.operand, names))
    if isinstance(node, ast.Constant):
        return node.value
    if isinstance(node, ast.Name):
        return names[node.id]
    if isinstance(node, ast.Call):
        return _FUNCS[node.func.id](*[_eval(a, names) for a in node.args])
    raise FormulaError(f"expression not allowed: {type(node).__name__}")


def compile_formula(expr, allowed_names):
    """Return fn(names_dict) -> number. Raises FormulaError on disallowed input."""
    try:
        tree = ast.parse(expr, mode="eval")
    except SyntaxError as e:
        raise FormulaError(f"bad formula {expr!r}: {e}")
    allowed = set(allowed_names)
    _validate(tree, allowed)

    def fn(names):
        return _eval(tree, names)

    fn.expr = expr
    return fn