diff options
author | Stas Medvedev <medvedevsa97@gmail.com> | 2024-06-12 16:56:07 +0300 |
---|---|---|
committer | Stas Medvedev <medvedevsa97@gmail.com> | 2024-06-12 16:56:07 +0300 |
commit | 0c1a65570be7f34f12a35da45669676f4479abd4 (patch) | |
tree | eb0905e4a6454d154b679830666e31f066a9fea7 /utils/restricted_exec.py | |
parent | ed49bb17b9e93a1406ab51f7dca5906661863627 (diff) |
add utils.restricted_exec
add app.tasks
Diffstat (limited to 'utils/restricted_exec.py')
-rw-r--r-- | utils/restricted_exec.py | 89 |
1 files changed, 89 insertions, 0 deletions
diff --git a/utils/restricted_exec.py b/utils/restricted_exec.py new file mode 100644 index 0000000..bf5205f --- /dev/null +++ b/utils/restricted_exec.py @@ -0,0 +1,89 @@ +import ast +import builtins +from copy import copy + +from RestrictedPython import compile_restricted_exec +from RestrictedPython.PrintCollector import PrintCollector +from RestrictedPython import safe_builtins +from RestrictedPython.Eval import ( + default_guarded_getiter, + default_guarded_getitem, + default_guarded_getattr +) +from RestrictedPython.Guards import guarded_iter_unpack_sequence + + +SAFE_MODULES = [ + 'math', + 'random', + 'math', + 'random', + 'time', + 'itertools', + 'functools', + 'operator', + 'collections', + 're', + 'json', + 'decimal', +] + +RESTRICTED_GLOBALS = { + '__builtins__': { + **safe_builtins.copy(), + '__import__': builtins.__import__ + }, + '_print_': PrintCollector, + '_getattr_': default_guarded_getattr, + '_getitem_': default_guarded_getitem, + '_getiter_': default_guarded_getiter, + '_iter_unpack_sequence_': guarded_iter_unpack_sequence, +} + + +def _restricted_exec(source, glb): + result = compile_restricted_exec(source) + assert result.errors == (), result.errors + assert result.code is not None + exec(result.code, glb) + return glb + + +def validate_code(code) -> list[str]: + err = [] + tree = ast.parse(code) + for node in ast.walk(tree): + if isinstance(node, (ast.Import, ast.ImportFrom)): + for node_name_obj in node.names: + if node_name_obj.name not in SAFE_MODULES: + msg = 'not allowed, %s' + if isinstance(node, ast.Import): + msg = msg % ('import %s' % node_name_obj.name) + else: + msg = msg % ('from %s import %s' % + (node.module, node_name_obj.name)) + + err.append(msg) + + return err + + +def restricted_exec(code) -> list[str]: + glb = _restricted_exec(code, copy(RESTRICTED_GLOBALS)) + + if '_print' in glb: + return glb['_print'].txt + else: + return [] + + +def getoutput(code: str) -> list[str]: + err = [] + try: + err.extend(validate_code(code)) + return restricted_exec(code) + + except Exception as error: + err.append('%s: %s' % (type(error).__name__, error)) + + return ['execution errors:', *err] |