app/services/py_sandbox.py

#!/usr/bin/env python3
import sys
import io
import os
import math as _math
import json as _json
import re as _re
import random as _random
import collections as _collections
import itertools as _itertools
import functools as _functools
import operator as _operator
import string as _string
import decimal as _decimal
import fractions as _fractions
import statistics as _statistics
import heapq as _heapq
import bisect as _bisect
import copy as _copy
import typing as _typing
import textwrap as _textwrap
import datetime as _datetime
import hashlib as _hashlib
import uuid as _uuid
import secrets as _secrets
import enum as _enum
import base64 as _base64
import time as _time

_SAFE_MODULES = {
    'math': _math, 'json': _json, 're': _re, 'random': _random,
    'collections': _collections, 'itertools': _itertools,
    'functools': _functools, 'operator': _operator, 'string': _string,
    'decimal': _decimal, 'fractions': _fractions, 'statistics': _statistics,
    'heapq': _heapq, 'bisect': _bisect, 'copy': _copy, 'typing': _typing,
    'textwrap': _textwrap, 'datetime': _datetime, 'hashlib': _hashlib,
    'uuid': _uuid, 'secrets': _secrets, 'enum': _enum, 'base64': _base64,
    'time': _time,
}

_ALLOWED_READ_DIRS = {os.path.realpath(p) for p in ('/tmp', os.getcwd()) if os.path.isdir(p)}


def _safe_import(name, *args, **kwargs):
    if name in _SAFE_MODULES:
        return _SAFE_MODULES[name]
    raise ImportError(f"Module '{name}' is not allowed in sandbox")


def _safe_open(file, mode='r', *args, **kwargs):
    if 'w' in mode or 'a' in mode or '+' in mode or 'x' in mode:
        raise PermissionError("File writing is not allowed in sandbox")
    real = os.path.realpath(os.path.abspath(str(file)))
    allowed = False
    for d in _ALLOWED_READ_DIRS:
        if real.startswith(d + os.sep) or real == d:
            allowed = True
            break
    if not allowed:
        raise PermissionError(f"File '{file}' is outside allowed read directories")
    if 'b' in mode:
        return io.open(file, mode, *args, **kwargs)
    return io.open(file, mode, *args, **kwargs)


def _safe_getattr(obj, name, *args):
    if isinstance(obj, (type, object)) and name in ('__class__', '__bases__', '__subclasses__',
                                                      '__mro__', '__globals__', '__code__',
                                                      '__closure__', '__dict__', '__builtins__'):
        raise AttributeError(f"Access to '{name}' is not allowed in sandbox")
    if args:
        return getattr(obj, name, args[0])
    return getattr(obj, name)


def _safe_setattr(obj, name, value):
    if name.startswith('_'):
        raise AttributeError(f"Setting attribute '{name}' is not allowed in sandbox")
    setattr(obj, name, value)


def _safe_delattr(obj, name):
    if name.startswith('_'):
        raise AttributeError(f"Deleting attribute '{name}' is not allowed in sandbox")
    delattr(obj, name)


_SAFE_BUILTINS = {
    '__import__': _safe_import,
    'print': print, 'len': len, 'range': range,
    'int': int, 'float': float, 'str': str, 'bool': bool,
    'list': list, 'dict': dict, 'tuple': tuple, 'set': set, 'frozenset': frozenset,
    'bytes': bytes, 'bytearray': bytearray,
    'True': True, 'False': False, 'None': None,
    'abs': abs, 'min': min, 'max': max, 'sum': sum, 'pow': pow,
    'round': round, 'divmod': divmod,
    'enumerate': enumerate, 'filter': filter, 'map': map, 'zip': zip,
    'all': all, 'any': any, 'iter': iter, 'next': next,
    'reversed': reversed, 'sorted': sorted, 'slice': slice,
    'bin': bin, 'chr': chr, 'hex': hex, 'oct': oct, 'ord': ord,
    'repr': repr, 'format': format, 'hash': hash, 'id': id,
    'type': type, 'isinstance': isinstance, 'issubclass': issubclass,
    'hasattr': hasattr, 'getattr': _safe_getattr, 'setattr': _safe_setattr,
    'delattr': _safe_delattr, 'callable': callable,
    'staticmethod': staticmethod, 'classmethod': classmethod,
    'property': property, 'object': object, 'super': super,
    'Exception': Exception, 'ValueError': ValueError, 'TypeError': TypeError,
    'KeyError': KeyError, 'IndexError': IndexError, 'AttributeError': AttributeError,
    'StopIteration': StopIteration, 'RuntimeError': RuntimeError,
    'ZeroDivisionError': ZeroDivisionError, 'ArithmeticError': ArithmeticError,
    'EOFError': EOFError, 'NameError': NameError, 'MemoryError': MemoryError,
    'NotImplementedError': NotImplementedError, 'OverflowError': OverflowError,
    'RecursionError': RecursionError, 'AssertionError': AssertionError,
    'ImportError': ImportError, 'PermissionError': PermissionError,
    'LookupError': LookupError, 'FileNotFoundError': FileNotFoundError,
    'input': input, 'open': _safe_open,
}

_FROZEN_ERR = TypeError("Cannot modify sandbox builtins")


class _FrozenDict(dict):
    def __setitem__(self, key, value):
        raise _FROZEN_ERR
    def __delitem__(self, key):
        raise _FROZEN_ERR
    def pop(self, key, *args):
        raise _FROZEN_ERR
    def popitem(self):
        raise _FROZEN_ERR
    def clear(self):
        raise _FROZEN_ERR
    def update(self, *args, **kwargs):
        raise _FROZEN_ERR
    def setdefault(self, key, *args):
        raise _FROZEN_ERR


def execute_code(code: str) -> None:
    try:
        compiled = compile(code, '<sandbox>', 'exec')
        globs = {'__builtins__': _FrozenDict(_SAFE_BUILTINS)}
        exec(compiled, globs)
    except Exception as e:
        print(str(e), file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    if len(sys.argv) > 1:
        with open(sys.argv[1], 'r') as f:
            code = f.read()
        execute_code(code)
    else:
        code = sys.stdin.read()
        execute_code(code)