Python 已经内置了针对 AST 的工具集:ast,可以在 Python 中完成与底层同样效果的工作。我们不仅可以通过 ast 库将任意代码编译为 AST 结构,而且对应的 AST 对象也可以被内置的 compile 函数编译为 CodeObject,并通过 exec 执行,这个过程也同样发生在 Python 底层。
# 在 Python 交互模式下
>>> import ast
>>> src='print("hello world.")'
>>> tree = ast.parse(src)
>>> co = compile(tree, "log.log", mode="exec")
>>> exec(co)
hello world.
代码被编译为 tree,然后被编译为 CodeObject,最终被执行,与直接执行 Python 源码等效。
准备工作,从文件读取源代码
只要有中间环节,那么就有机会在环节之间做一些小动作,比如在已经生成的 AST 中插入一些不存在的节点!
当我们修改了 AST,对于其后的环节并知道 AST 已经与最初的源码发生了变化,只要符合规则,VM 就会执行被修改的 AST,这样就 Hack 了 AST!思路非常清晰,下面先做一点准备工作。
让我们对上面的例程做一些调整,改为从一个独立的 target.py 中读取代码。
# main.py
import ast
# 从文件中读取代码
with open("target.py") as f:
tree = ast.parse(f.read())
import astpretty
# 打印 AST
astpretty.pprint(tree)
# 编译,执行 AST
co = compile(tree, "log.log", mode="exec")
exec(co)
与之对应的 target.py 的内容是这样的:
def test():
return "Go!!!"
print("Hey girl.", test())
这次,我们会解析更高级的 AST 结构:
- 一个函数定义;
- 和一个函数调用,以及在参数中调用另一个函数。
只要运行:
python main.py
输出非常长,可以迅速下拉:
Module(
body=[
FunctionDef(
lineno=1,
col_offset=0,
end_lineno=2,
end_col_offset=18,
name='test',
args=arguments(posonlyargs=[], args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]),
body=[
Return(
lineno=2,
col_offset=4,
end_lineno=2,
end_col_offset=18,
value=Constant(lineno=2, col_offset=11, end_lineno=2, end_col_offset=18, value='Go!!!', kind=None),
),
],
decorator_list=[],
returns=None,
type_comment=None,
),
Expr(
lineno=4,
col_offset=0,
end_lineno=4,
end_col_offset=26,
value=Call(
lineno=4,
col_offset=0,
end_lineno=4,
end_col_offset=26,
func=Name(lineno=4, col_offset=0, end_lineno=4, end_col_offset=5, id='print', ctx=Load()),
args=[
Constant(lineno=4, col_offset=6, end_lineno=4, end_col_offset=17, value='Hey girl.', kind=None),
Call(
lineno=4,
col_offset=19,
end_lineno=4,
end_col_offset=25,
func=Name(lineno=4, col_offset=19, end_lineno=4, end_col_offset=23, id='test', ctx=Load()),
args=[],
keywords=[],
),
],
keywords=[],
),
),
],
type_ignores=[],
)
Hey girl. Go!!!
最终的结果相当于直接执行 targe.py。注意观察 AST 的结构,test 在 print 函数的 args[1] 位置,与源码中 test 函数在参数列表中的位置一致,这也可以看出 AST 与源码之间的对应关系。
Hack AST
下面我们通过 ast 库提供的辅助类,将 print 函数的第一个参数修由 "Hey girl." 改为经典的 “hello world”。
基本思路是:
- 创建一个新的 Constant 类节点,包含 “hello world.”;
- 用这个新节点取代 args[0];
- 编译、执行。
AST 是一个比较复杂的结构,ast 库提供了 ast.NodeVisitor 类,它通常作为节点访问类的基类。具体使用方法可以参考官方文档 和更加详实的开源文档。
# main.py
import ast
# 从文件中读取代码
with open("target.py") as f:
tree = ast.parse(f.read())
# 修改前的结果
print("origin:")
co = compile(tree, "log.log", "exec")
exec(co)
# 构建一个节点访问器,对 Call 类节点进行操作,类名称可以从 AST 输出看到就是 Call
class Visitor(ast.NodeVisitor):
def visit_Call(self, node):
# 只处理 print 函数
if node.func.id != "print":
return
print("set arg to 'hello world'")
# 用新的 Constant 节点覆盖 args[0]
args = node.args
args[0] = ast.Constant(s = "hello world", kind=None)
# AST 节点的行号信息必须填写,该函数将使用父节点的行号信息填充当前节点。
ast.fix_missing_locations(args[0])
# 创建访问器
v = Visitor()
# 访问 AST, 并对 print 函数第一个参数修改
v.visit(tree)
# 修改后
print("hacked AST:")
# 编译执行
co = compile(tree, "log.log", "exec")
exec(co)
输出:
origin:
Hey girl. Go!!!
set arg to 'hello world'
hacked AST:
hello world Go!!!
Look,成功了,我们只是控制了 AST 就使得代码运行结果与源码不同了!
从 AST 到源码
Python 生态提供了很多有用的工具,AST 与源码的逻辑关系如此紧密,可以轻松的从 AST 转回 Python 源码,可以通过 astunparse 库实现。
# 接上面的代码
import astunparse
print("modified code:")
print(astunparse.unparse(tree))
输出:
modified code:
def test():
return 'Go!!!'
print('hello world', test())
我们对 AST 所做的一切,相当于修改了 target.py 中 print 函数的第一个参数。
AST 的用处
AST 相比源码有更高的结构性和抽象性,可以从更底层处理代码。比如在 AST 中可以检测到特定的函数定义、包导入、函数调用等等,如果有禁用或者限制、增强某些调用的需要,就可以对源码的 AST 进行调整后,再交给 VM 执行得到预期的效果。如果没有 AST,很多修改可能就需要复杂且不可靠的源码字符串操作了。
AST 可以看作是结构化的代码,在代码的静态分析中比较常用,有比直接面向源码更高的性能和便利性。如果有必要,甚至可以根据 AST 自己实现新的 Python 解释器,或者对 AST 做一些优化。Python 提供的 ast 工具包非常有诚意,Python 实现的细节一览无余,阅读 AST 比直接阅读 ByteCode 更加友好。