python ast代码检查，禁止导入危险函数包

import ast


class CheckFun(ast.NodeVisitor):
    def __init__(self):
        super().__init__()
        self.ban_moudel = [
            "os",
            "sys",
            "socket",
            "multiprocessing",
            "requests",
        ]
        self.ban_func = ["exec", "eval"]
        self.allow = True
        self.res = []

    def visit_Import(self, node):
        # 防止导入包
        for item in node.names:
            if item.name in self.ban_moudel:
                self.allow = False
                self.res.append("不允许导入 %s包" % item.name)

    def visit_Assign(self, node):
        # 防止赋值后调用
        if node.value.id in self.ban_func:
            self.allow = False
            self.res.append("赋值右边不允许写 %s " % node.value.id)

    def visit_Call(self, node):
        # 防止直接调用
        if isinstance(node, ast.Name):
            if node.func.id in self.ban_func:
                self.allow = False
                self.res.append("不允许调用 %s " % node.func.id)

    def visit_ImportFrom(self, node):
        # 防止直接调用
        print(node.module)
        if node.module in self.ban_moudel:
            self.allow = False
            self.res.append("不允许导入 %s包" % node.module)


def test_function(fn_str):
    root_node = ast.parse(fn_str)
    ckf = CheckFun()
    ckf.visit(root_node)
    if ckf.allow:
        print("allow")
    else:
        print("unallow")
        for item in ckf.res:
            print(item)


func = """
from requests import xx
import socket
import os
import sys
b=eval
b(123)
eval(321)
def aaaaaa(text):
    import sx
    import sys
    print(1)
    return re.search(r"\d+", text).group(0)
"""

test_function(func)