基于Python AST實現(xiàn)代碼安全檢測功能

更新時間：2025年12月03日 09:05:34 作者：weixin_46244623

在某些場景下,我們必須限制用戶代碼中危險模塊,所以本文將和大家介紹一下基于 Python AST的代碼安全檢查方案,并給出完整示例代碼,有需要的可以了解下

一、為什么使用 AST

常見的字符串匹配方式容易被繞過，例如：

e = ev
b = al
(e + b)("print(123)")

而 AST 解析后，所有語句都會被解析成結(jié)構(gòu)化節(jié)點，從根本上杜絕字符串拼接繞過的可能性。

二、完整檢測代碼

下面是完整的 AST 檢測類 CheckFun：

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):
        # 禁止賦值后調(diào)用危險函數(shù)，如 b = eval
        try:
            if node.value.id in self.ban_func:
                self.allow = False
                self.res.append("賦值右側(cè)不允許出現(xiàn) %s" % node.value.id)
        except:
            pass

    def visit_Call(self, node):
        # 禁止直接調(diào)用 eval、exec
        try:
            if isinstance(node.func, ast.Name):
                if node.func.id in self.ban_func:
                    self.allow = False
                    self.res.append("不允許調(diào)用 %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)

三、測試示例

示例一：包含危險模塊和危險函數(shù)的代碼

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)
"""

檢測執(zhí)行：

test_function(func2)

輸出示例：

e = ev
b = al
(e + b)("print(123)")

四、從文件讀取并檢測

如果需要從外部文件讀取用戶代碼并檢測，可以這樣寫：

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")