Python AST 研究
Preface
抽象语法树(AST)，即Abstract Syntax Tree的缩写。它是源代码语法结构的一种抽象表示。它以树状的形式表现编程语言的语法结构，树上的每个节点都表示源代码中的一种结构。之所以说语法是抽象的，是因为这里的语法并不会表示出真实语法中出现的每个细节。
Python内置一个ast库，其中存在一些内置方法生成，遍历，获取，解析 Python代码并获取ast树，但ast库的官方文档过于简陋，基本上无法当作学习文档。
最简单的Demo
import ast
import astpretty
res = ast.parse("a = 1+1")
astpretty.pprint(res)Module(
body=[
Assign(
lineno=1,
col_offset=0,
end_lineno=1,
end_col_offset=7,
targets=[Name(lineno=1, col_offset=0, end_lineno=1, end_col_offset=1, id='a', ctx=Store())],
value=BinOp(
lineno=1,
col_offset=4,
end_lineno=1,
end_col_offset=7,
left=Constant(lineno=1, col_offset=4, end_lineno=1, end_col_offset=5, value=1, kind=None),
op=Add(),
right=Constant(lineno=1, col_offset=6, end_lineno=1, end_col_offset=7, value=1, kind=None),
),
type_comment=None,
),
],
type_ignores=[],
)
Process finished with exit code 0astpretty是一个第三方库，用来美化输出ast树
上面的输出结果就是，a = 1 + 1 这句语句对应的ast树的可读形式，
其中Module(body(是所有语句dump出来都会有的，不用管
由于a = 1 + 1是一句赋值语句，所以第一层为 Assign语法
对于每一种结构语法，都会自己对应的语法属性，比如这里的Assign，就存在targets(赋值对象),value(被赋的值),而这些语法属性自己又包含了自己的属性，一层一层嵌套
需要注意的是对于某些表达式，python的ast的遍历顺序是从后到前，从外到内，这个与php的ast遍历顺序不同,有点反人类
比如
request.arg.get(123)
func(a[:5])
第一个获取到的ast对象是 get(123)
第二个获取到的也是函数 func()
当然除了打印ast树，我们也可以直接输出ast树对应的节点
假设有一道CTF题目
#secret.py
def func():
flag = "xxxxxxxxxxxxxxxxxxxxxxxxx"
# guess flag
假如我们可以运行一段py脚本,但是不允许读文件,执行系统命令，我们该如何获取flag?
#user.py
import secret
最简单的办法是py2的co_constco_const属性，可以获得一个函数定义块中被定义的常量值
ast增加节点
那如果是Py3 或者 secret.py中的flag不以常量方式赋值怎么办呢
拿下面的secret.py举个例子，假设我们需要想办法获取flag的值
#secret.py
import uuid
def check():
# Example
flag = f"flag{{{uuid.uuid4()}}}"
现在我们允许读取文件，但是我们无法获取到uuid的值
我们可以修改ast树，在赋值的下面加一句print的ast树，然后调用ast的eval运行此函数
首先查看一个pring(flag)的ast节点
import ast
import astpretty
res = ast.parse("print(flag)")
astpretty.pprint(res)Module(
body=[
Expr(
lineno=1,
col_offset=0,
end_lineno=1,
end_col_offset=11,
value=Call(
lineno=1,
col_offset=0,
end_lineno=1,
end_col_offset=11,
func=Name(lineno=1, col_offset=0, end_lineno=1, end_col_offset=5, id='print', ctx=Load()),
args=[Name(lineno=1, col_offset=6, end_lineno=1, end_col_offset=10, id='flag', ctx=Load())],
keywords=[],
),
),
],
type_ignores=[],
)
我们模仿这个AST树，用python代码生成对应的ast结构
import ast
func = ast.Name(id="print", ctx=ast.Load())
args = [ast.Name(id="flag", ctx=ast.Load())]
call = ast.Call(func=func, args=args, keywords=[])
node = ast.Expr(value=call)然后把我们自己生成的节点，插入到check函数中
content = open("secret.py").read()
res = ast.parse(content)
res.body[1].body.insert(1,node)
ast.fix_missing_locations(res)
由于print(flag)与赋值语句在同一层
而res.body[1]是函数定义结构，res.body[1].body即为函数体内，所以在此处插入res.body[1].body[0]是赋值语句，所以为insert(1,node)
ast.fix_missing_locations函数作用是自动修复 ast结构的偏移 行号等杂项，需要调用，否则会报
TypeError: required field "lineno" missing from stmt
最后只需要编译运行我们的ast树即可
exec(compile(res, '', 'exec'))
由于我们运行的语句相当于重新定义了一个check函数，所以还需要运行一遍check函数
整个文件结构如下
import ast
content = open("secret.py").read()
func = ast.Name(id="print", ctx=ast.Load())
args = [ast.Name(id="flag", ctx=ast.Load())]
call = ast.Call(func=func, args=args, keywords=[])
node = ast.Expr(value=call)
res = ast.parse(content)
res.body[1].body.insert(1,node)
ast.fix_missing_locations(res)
exec(compile(res, '', 'exec'))
check()ast修改节点
除了可以在下一行新增节点以外，我们还可以直接修改赋值节点，直接改为表达式节点将值输出
跟上面差不多，就不多介绍了
import ast
content = open("secret.py").read()
res = ast.parse(content)
value = res.body[1].body[0].value
func = ast.Name(id="print", ctx=ast.Load())
args = [value]
call = ast.Call(func=func, args=args, keywords=[])
node = ast.Expr(value=call)
res.body[1].body[0] = node
ast.fix_missing_locations(res)
exec(compile(res, '', 'exec'))
check()
res.body[1].body[0].value 就是赋值表达式所赋的值
ast删除节点
body是一个list，把对应节点pop就行了
#secret.py
import uuid
def check():
flag = f"flag{{{uuid.uuid4()}}}"
flag = "88888"
print(flag)
import ast
import astpretty
content = open("secret.py").read()
func = ast.Name(id="print", ctx=ast.Load())
args = [ast.Name(id="flag", ctx=ast.Load())]
call = ast.Call(func=func, args=args, keywords=[])
node = ast.Expr(value=call)
res = ast.parse(content)
res.body[1].body.pop(1)
ast.fix_missing_locations(res)
astpretty.pprint(res)
exec(compile(res, '', 'exec'))
check()遍历节点
以上内容只是对某单一节点进行修改，如果我们要修改所有某一特征节点时如何去做。
ast库提供了 visit方法，供开发者遍历所有节点
import ast
content = open("secret.py").read()
res = ast.parse(content)
class MyVisit(ast.NodeVisitor):
def visit_Assign(self,node):
print(node.value)
return node
m = MyVisit()
m.visit(res)<_ast.JoinedStr object at 0x02F2F6B8>
<_ast.Constant object at 0x02F2F928>
开发者需要定义一个ast.NodeVisitor的子类，然后定义对应的visit_结构类型方法，处理特定的结构。
然后创建这个类对象，调用visit方法。遍历AST树
但此类只能进行ast的输出，如果需要一边遍历一边修改ast，增删改等，就需要继承另一个子类ast.NodeTransformer
假如我们把所有常量数字改为字符串
我们先看看字符串和数字的ast结构
import ast
import astpretty
res = ast.parse("a = 123\nv='123'")
astpretty.pprint(res)Module(
body=[
Assign(
lineno=1,
col_offset=0,
end_lineno=1,
end_col_offset=7,
targets=[Name(lineno=1, col_offset=0, end_lineno=1, end_col_offset=1, id='a', ctx=Store())],
value=Constant(lineno=1, col_offset=4, end_lineno=1, end_col_offset=7, value=123, kind=None),
type_comment=None,
),
Assign(
lineno=2,
col_offset=0,
end_lineno=2,
end_col_offset=7,
targets=[Name(lineno=2, col_offset=0, end_lineno=2, end_col_offset=1, id='v', ctx=Store())],
value=Constant(lineno=2, col_offset=2, end_lineno=2, end_col_offset=7, value='123', kind=None),
type_comment=None,
),
],
type_ignores=[],
)
Process finished with exit code 0
很明显，他们都属于Constant结构，完全对应里面的Value属性
需要注意的是，在老的python版本里，这些常量值分别属于ast.Num, ast.Str而不是ast.Constant
所以这个数字转字符串的遍历类如下
#case.py
a = 123
b = "231445"
def s():
c= 3124235213#exp.py
import ast
import astunparse
content = open("case.py").read()
res = ast.parse(content)
class MyVisit(ast.NodeTransformer):
def visit_Constant(self,node):
if type(node.value) == int:
node.value = str(node.value)
return node
m = MyVisit()
node = m.visit(res)
print(astunparse.unparse(node))
a = '123'
b = '231445'
def s():
c = '3124235213'
Process finished with exit code 0astunparse是一个三方库，将ast树转化为python code
而如果需要删除ast节点就更简单了，在方法里return None即可
Introductionas
https://stackoverflow.com/questions/46388130/insert-a-node-into-an-abstract-syntax-tree
https://www.escapelife.site/posts/40a2fc93.html
