基于PythonAST实现代码安全检测功能

目录

• 一、为什么使用 AST
• 二、完整检测代码
• 三、测试示例
• 四、从文件读取并检测
• 五、实际应用场景

在某些场景下（在线代码执行平台、脚本上传平台、自动任务系统等），我们必须限制用户代码中危险模块（如 os、sys、socket）以及危险函数（如 eval、exec）的使用，否则会带来安全风险。

本文介绍一种 基于 Python AST（抽象语法树） 的代码安全检查方案，并给出完整示例代码，适用于实际生产环境。

常见的字符串匹配方式容易被绕过，例如：

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

而 AST 解析后，所有语句都会被解析成结构化节点，从根本上杜绝字符串拼接绕过的可能性。

下面是完整的 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):
        # 禁止赋值后调用危险函数，如 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)

输出示例：

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

该方式非常适合：

• 在线代码执行平台
• WebIDE / 教程平台
• 自动 Python 任务系统
• 安全沙箱
• 禁止用户访问系统资源的场景

因为 AST 是结构化的，无法通过字符串拼接绕过，是目前最可靠的静态检测方式之一。

到此这篇关于基于Python AST实现代码安全检测功能的文章就介绍到这了,更多相关Python代码检测内容请搜索风君子博客以前的文章或继续浏览下面的相关文章希望大家以后多多支持风君子博客！

您可能感兴趣的文章:

• Python实现异常检测LOF算法的示例代码
• 深入浅析Python代码规范性检测
• Python AST 模块实战演示
• python ast模块详析与用法