基于Python AST实现代码安全检测功能
作者:weixin_46244623
在某些场景下,我们必须限制用户代码中危险模块,所以本文将和大家介绍一下基于 Python AST的代码安全检查方案,并给出完整示例代码,有需要的可以了解下
在某些场景下(在线代码执行平台、脚本上传平台、自动任务系统等),我们必须限制用户代码中危险模块(如 os、sys、socket)以及危险函数(如 eval、exec)的使用,否则会带来安全风险。
本文介绍一种 基于 Python AST(抽象语法树) 的代码安全检查方案,并给出完整示例代码,适用于实际生产环境。
一、为什么使用 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代码检测内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家!
