12月28, 2020

59. Hack 抽象语法树

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 中插入一些不存在的节点!

image.png

当我们修改了 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”。

image.png

基本思路是:

  • 创建一个新的 Constant 类节点,包含 “hello world.”;
  • 用这个新节点取代 args[0];
  • 编译、执行。

AST 是一个比较复杂的结构,ast 库提供了 ast.NodeVisitor 类,它通常作为节点访问类的基类。具体使用方法可以参考官方文档 和更加详实的开源文档

image.png

# 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 更加友好。

本文链接:http://www.thinkinpython.com/post/deep_python_vm_59.html

-- EOF --