快捷導(dǎo)航

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

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

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

在某些場(chǎng)景下（在線代碼執(zhí)行平臺(tái)、腳本上傳平臺(tái)、自動(dòng)任務(wù)系統(tǒng)等），我們必須限制用戶代碼中危險(xiǎn)模塊（如 os、sys、socket）以及危險(xiǎn)函數(shù)（如 eval、exec）的使用，否則會(huì)帶來(lái)安全風(fēng)險(xiǎn)。

本文介紹一種 基于 Python AST（抽象語(yǔ)法樹(shù)） 的代碼安全檢查方案，并給出完整示例代碼，適用于實(shí)際生產(chǎn)環(huán)境。

一、為什么使用 AST

常見(jiàn)的字符串匹配方式容易被繞過(guò)，例如：

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

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

二、完整檢測(cè)代碼

下面是完整的 AST 檢測(cè)類 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):
        # 禁止導(dǎo)入模塊
        for item in node.names:
            if item.name in self.ban_moudel:
                self.allow = False
                self.res.append("不允許導(dǎo)入 %s 包" % item.name)

    def visit_Assign(self, node):
        # 禁止賦值后調(diào)用危險(xiǎn)函數(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("不允許導(dǎo)入 %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)

三、測(cè)試示例

示例一：包含危險(xiǎn)模塊和危險(xiǎn)函數(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)
"""

檢測(cè)執(zhí)行：

test_function(func2)

輸出示例：

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

四、從文件讀取并檢測(cè)

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

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