基于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):
        # 禁止赋值后调用危险函数，如 b = eval
        try:
            if node.value.id in self.ban_func:
                self.allow = False
                self.res.append("赋值右侧不允许出现 %s" % node.value.id)
        except:
            pass

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

    def visit_ImportFrom(self, node):
        # 禁止 from xxx import xx
        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)

func2 = """
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
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(func2)

with open('./tmp/' + args[1], 'r') as f:
    root_node = ast.parse(f.read())
    ckf = CheckFun()
    ckf.visit(root_node)
    if ckf.allow:
        print("allow")
    else:
        print("unallow")