codesafety.py

import ast
import sys
from collections import defaultdict

BANNED_NAMES = [
    'exec', 'eval', 'compile', 'open', 'input', '__import__'
    , 'sys', 'subprocess'
]

BANNED_NODES = [
    ast.Import, ast.ImportFrom, ast.Global, ast.Nonlocal
]

def codeid(fname, code):
    return f"{fname} -> `{code[:50]}`"

class MainVisitor(ast.NodeVisitor):
    """
    Walks the AST of a user-defined main function to identify
    if uses of BANNED_NAMES exist or BANNED_NODES are part of
    the tree
    """

    def __init__(self):
        self.errors = []

    def _error_dict(self):
        d = defaultdict(list)

        for id, data in self.errors:
            d[id].append(data)

        return d

    def visit(self, node):
        if type(node) in BANNED_NODES:
            self.errors.append(('Type Banned', type(node)))

        return super().visit(node)

    def visit_Name(self, node):
        if node.id in BANNED_NAMES:
            self.errors.append(('Name Banned', node.id))

        self.generic_visit(node)

    def visit_Attribute(self, node):
        if node.attr.startswith('__') and node.attr.endswith('__'):
            self.errors.append(('Dunder Access', node.attr))

        self.generic_visit(node)

class CodeSafety:
    """
    Verifies a user-defined script is "safe" to a certain degree

    Class Attributes:
        CODE_BYTE_LIMIT : int
            max size in bytes allowed for the user-defined function
    """

    CODE_BYTE_LIMIT = 1500

    @classmethod
    def check_user_main_function(cls, fname, code):
        if sys.getsizeof(code) > cls.CODE_BYTE_LIMIT:
            raise Exception(f"{codeid(fname, code)} exceeds byte limit {cls.CODE_BYTE_LIMIT}!")

        tree = ast.parse(code)

        if len(tree.body) != 1:
            raise Exception(f"{codeid(fname, code)} is not a single code block and contains multiple (or 0) nodes!")

        if not isinstance(tree.body[0], ast.FunctionDef):
            raise Exception(f"{codeid(fname, code)} does not start with a function definition")

        if tree.body[0].name != 'main':
            raise Exception(f"{codeid(fname, code)} function name is not `main`!")

        visitor = MainVisitor()
        visitor.visit(tree)

        if visitor.errors:
            raise Exception(f"{codeid(fname, code)} has AST unsafe errors:\n{visitor._error_dict()}")

        return True