Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py: 69%

1""" 

2RestrictedPython-based sandbox for custom code guardrails. 

3 

4User-supplied guardrail source is compiled with ``compile_restricted`` and 

5executed with curated globals. All attribute access, subscripting, iteration, 

6and assignment is mediated by RestrictedPython guards, which block dunder 

7access (``__globals__``, ``__code__``, ``__class__``, ``__setattr__``, etc.) 

8and reject dangerous AST constructs (``import``, ``exec``, ``eval``, 

9``compile``, class definitions, etc.) at compile time. 

10 

11The default ``RestrictingNodeTransformer`` denies ``async def``/``await``, 

12which breaks the documented async guardrail pattern (``await http_get(...)``). 

13We subclass it to permit those specific nodes, while keeping every other 

14restriction intact. 

15""" 

16 

17import ast 

18import operator 

19from collections.abc import Callable, Mapping 

20from types import CodeType 

21from typing import Final 

22 

23from RestrictedPython import ( 

24 RestrictingNodeTransformer, 

25 compile_restricted, 

26 limited_builtins, 

27 safe_builtins, 

28 utility_builtins, 

29) 

30from RestrictedPython.Eval import default_guarded_getitem, default_guarded_getiter 

31from RestrictedPython.Guards import ( 

32 full_write_guard, 

33 guarded_iter_unpack_sequence, 

34 safer_getattr, 

35) 

36 

37from .primitives import get_custom_code_primitives 

38 

39 

40class AsyncAwareTransformer(RestrictingNodeTransformer): 

41 """Extend the default transformer to allow ``async def`` and ``await``. 

42 

43 The base class rejects every async AST node outright. ``AsyncFunctionDef`` 

44 has the same ``_fields`` as ``FunctionDef`` and the same security 

45 semantics, so we delegate to ``visit_FunctionDef`` — name check, argument 

46 check, print-scope wrapping, and any future additions to that method are 

47 inherited automatically. ``AsyncFor``/``AsyncWith``/``Await`` delegate to 

48 ``node_contents_visit`` so their children still get transformed. 

49 """ 

50 

51 def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AST: 

52 return self.visit_FunctionDef(node) 

53 

54 def visit_AsyncFor(self, node: ast.AsyncFor) -> ast.AST: 

55 return self.node_contents_visit(node) 

56 

57 def visit_AsyncWith(self, node: ast.AsyncWith) -> ast.AST: 

58 return self.node_contents_visit(node) 

59 

60 def visit_Await(self, node: ast.Await) -> ast.AST: 

61 return self.node_contents_visit(node) 

62 

63 

64_INPLACE_OPS: Final[Mapping[str, Callable[[object, object], object]]] = { 

65 "+=": operator.iadd, 

66 "-=": operator.isub, 

67 "*=": operator.imul, 

68 "/=": operator.itruediv, 

69 "//=": operator.ifloordiv, 

70 "%=": operator.imod, 

71 "**=": operator.ipow, 

72 "@=": operator.imatmul, 

73 "&=": operator.iand, 

74 "|=": operator.ior, 

75 "^=": operator.ixor, 

76 "<<=": operator.ilshift, 

77 ">>=": operator.irshift, 

78} 

79 

80 

81def _inplacevar_(op: str, x: object, y: object) -> object: 

82 # RestrictedPython rewrites ``x += 1`` on a simple name into 

83 # ``x = _inplacevar_("+=", x, 1)``. The package deliberately ships no 

84 # default, so we dispatch through ``operator``'s in-place helpers, which 

85 # honour Python's normal ``__iadd__``/``__add__`` fallback. 

86 fn: Final = _INPLACE_OPS.get(op) 

87 if fn is None: 

88 raise SyntaxError(f"augmented assignment {op!r} is not supported") 

89 return fn(x, y) 

90 

91 

92def _build_sandbox_builtins() -> dict[str, object]: 

93 # ``limited_builtins`` overrides ``list``/``tuple``/``range`` from 

94 # ``safe_builtins`` with bounds-checking variants (e.g. ``limited_range`` 

95 # rejects ``range(10**18)``). ``utility_builtins`` adds ``set``, 

96 # ``frozenset``, ``math``, ``random``, and a filtered ``string`` delegator. 

97 return { 

98 **safe_builtins, 

99 **limited_builtins, 

100 **utility_builtins, 

101 } 

102 

103 

104def build_sandbox_globals() -> dict[str, object]: 

105 """Assemble the globals dict for executing guardrail code. 

106 

107 Includes the LiteLLM-provided primitives (``regex_match``, ``http_get``, 

108 ``allow``/``block``/``modify``, etc.) plus the RestrictedPython guards 

109 that the compiled bytecode expects to find by name. 

110 """ 

111 return { 

112 **get_custom_code_primitives(), 

113 "__builtins__": _build_sandbox_builtins(), 

114 "_getattr_": safer_getattr, 

115 "_getitem_": default_guarded_getitem, 

116 "_getiter_": default_guarded_getiter, 

117 "_iter_unpack_sequence_": guarded_iter_unpack_sequence, 

118 "_write_": full_write_guard, 

119 "_inplacevar_": _inplacevar_, 

120 } 

121 

122 

123def compile_sandboxed(source: str, filename: str = "<guardrail>") -> CodeType: 

124 """Compile guardrail source with RestrictedPython's AST transformer. 

125 

126 Raises ``SyntaxError`` on either a Python syntax error or a restricted 

127 construct (import, exec, dunder name, etc.). 

128 """ 

129 return compile_restricted( 

130 source=source, 

131 filename=filename, 

132 mode="exec", 

133 policy=AsyncAwareTransformer, 

134 )