weggli debug

工具

关于Weggli

AST Pattern Search

核心是使用和 tree-sitter 库，然后搞了 query-tree 来在 AST上进行搜索，这只能说是匹配特定的代码片段，还达不到程序分析的那个级别，所以理论上只能过程内分析，而且没有上下文啥的 :D 直白点说的话，像是AST的正则表达式，不过某种意义上来说对于使用白盒方案快速召回一些漏洞也是一种借鉴吧。

当然我也用这个工具做了一些扩展，结合其他工具解决了一些问题，目前看来这个东西还是具有一定的可玩性的 :D

Weggli如何工作

看代码，调试分析

idea配置

安装Rust插件，调试的话，会默认再去安装Native Debugging Support，有了这俩东西就可以调试了

配置传递给weggli的参数的话跟在 -- 后面即可 : 

1
run --package weggli --bin weggli --  "{$func($b);system($b);}" -R "func=printf$" /path/to/src

工作流程

只描述核心流程

query-tree 构建

参考 tree-sitter文档

1
2
3
4
5
6
7
8
9
10
let work: Vec<WorkItem> = args
        .pattern
        .iter()
        .map(|pattern| {
            match parse_search_pattern(pattern, args.cpp, args.force_query, &regex_constraints) {
                Ok(qt) => {
                    let identifiers = qt.identifiers();
                    variables.extend(qt.variables());
                    WorkItem { qt, identifiers }
        // ....

构造 WorkItem{qt, identifiers}

  • qt : query-tree, tree-sitter的Tree对象
  • identifiers : 标识符，query中”终结符”

调用链:

1
2
3
4
5
6
main
    parse_search_pattern
        weggli::parse(pattern, is_cpp) // 返回Tree对象
        //修正pattern
        validate_query
        build_query_tree

修正pattern : weggli处理了“不合法的”格式，如:

  • memcpy(a,b,size) -> memcpy(a,b,size);
  • memcpy(a,b,size); -> {memcpy(a,b,size);}
1
validate_query(&tree, p, force_query)? // 返回 TreeCursor对象，用于遍历AST

语法合法性检查，如果 force_query为True，意味着忽略这些语法错误

如 :

image-20220707154439731

1
"{$func($b);_($b);}"

对应 : 

1
2
3
4
5
6
7
8
9
10
11
(translation_unit 
    (
        compound_statement 
        (
            expression_statement (call_expression function: (identifier) arguments: (argument_list (identifier)))
        ) 
        (
            expression_statement (call_expression function: (identifier) arguments: (argument_list (identifier)))
        )
    )
)

同时还不允许 :

image-20220707155732813

返回的是 : c.goto_first_child();，即 花括号中间的内容

1
2
3
4
5
6
7
8
Ok(build_query_tree(
 p,
 &mut c,
 is_cpp,
 Some(regex_constraints.clone()),
))

_build_query_tree(source, cursor, 0, is_cpp, false, false, regex_constraints)

QueryTree数据结构:

1
2
3
4
5
6
7
pub struct QueryTree {
    query: Query,
    captures: Vec<Capture>,
    negations: Vec<NegativeQuery>,
    variables: HashSet<String>,
    id: usize,
}

转换的tree_sitter query (核心逻辑都在 builder.rs 的 QueryBuilder.build) 

1
2
Translate the tree below `c` into a tree-sitter query string.
"{$func($b);_($b);}"
1
2
tree_sitter query 1: ((call_expression function:[(identifier) (field_expression) (field_identifier)] @0 arguments:(argument_list [(identifier) (field_expression) (field_identifier)] @1)) )([(identifier) (field_expression) (field_identifier)] @2 )
tree_sitter query 0: (function_definition body: (compound_statement) @0) @1

深度优先的方式递归生成query tree string，按照AST解析出来不同的节点，后面跟着的 @x 用来区分不同的 identifier，方便后面做匹配。

如简单的 {printf(var, bar);} 生成的 query-tree是 : 

1
2
3
4
5
6
7
8
((call_expression         
        function:  [(field_expression field: (field_identifier)@0) (identifier) @0] 
        arguments: (argument_list 
                        . (identifier) @1 
                        . (identifier) @2)
                    ) 
                        
                        (#eq? @0 "printf")(#eq? @1 "var")(#eq? @2 "bar")) // captures

结合tree-sitter的playground来看就很容易看明白了:

tree-sitter-playground

query执行(pattern 匹配)

在执行query之前会做

  • 对于需要正则匹配的 identifer做合法性确认

    1
    2
    3
    4
    5
    6
    for v in regex_constraints.variables() {
            if !variables.contains(v) {
                eprintln!("'{}' is not a valid query variable", v.red());
                std::process::exit(1)
            }
        }

  • 确定待解析源码文件(Verify that the --include and --exclude regexes are valid.) 主要是根据后缀来

随后就是通过管道来处理，分为：

  • 文件读取 & AST解析 let (ast_tx, ast_rx) = mpsc::channel();
  • QueryTree 匹配 & 结果输出 let (results_tx, results_rx) = mpsc::channel();
1
2
3
4
5
6
7
8
9
10
11
12
// Spawn worker to iterate through files, parse potential matches and forward ASTs
s.spawn(move |_| parse_files_worker(files, ast_tx, w, cpp));

// Run search queries on ASTs and apply CLI constraints
// on the results. For single query executions, we can
// directly print any remaining matches. For multi
// query runs we forward them to our next worker function
s.spawn(move |_| execute_queries_worker(ast_rx, results_tx, w, &args));

if w.len() > 1 {
    s.spawn(move |_| multi_query_worker(results_rx, w.len(), before, after));
}

**这玩意描述起来就像个流水线 :D **

详细描述的话就是：在有了 query-tree就需要把目标文件，解析(parse_files_worker)成 (Tree, source_code)，结果发送到 ast_tx，然后从ast_rx获取这些信息来执行查询操作(execute_queries_worker)；结果放在 result_tx，后面处理结果的函数会从result_rx获取，然后输出。

1
2
3
4
5
6
7
parse_files_worker(files, ast_tx, w, cpp);
    weggli::parse(....);
execute_queries_worker(ast_rx, results_tx, w, &args); // w  WorkItem，里面有query-tree
    qt.matches(tree.root_node(), &source);
    match_internal(...);
    QueryCursor.matches(...);
    QueryTree.process_match(...);

TODO: 需要细读逻辑

这里简单的加一句print之类的可以来看看每次query的时候目标tree是啥样的(生成过程和query tree类似)

1
2
3
4
// Run query
let tmp_tree = tree.root_node().to_sexp();
let matches = qt.matches(tree.root_node(), &source);

image-20220801202121736

所以这就转换成了一个字符串匹配的问题，结合之前的 -R ，能支持正则匹配，所以说weggli是在AST上搞正则匹配一点都没说错 :D

multi-query(-p 参数)

漏洞模型测试

Question - query construction 这个issue里提到了这个场景，先还原一下场景 :

vuln.c 是个类似的情况，尝试query

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

int wtf(int a){
    return a + 1337;
}


int foo(int bar){

    wtf(bar);

    system(bar);
}


int vuln(char *data){

    char cmd[2048] = {0};

    sprintf(cmd, "/bin/bash %s > /tmp", data);

    return system(cmd);
}


int main(int argc, char*argv[]){


    if (argc < 2){
        return 0;
    }

    foo(11111);
    vuln(argv[1]);
    return 0;
}

image-20220803130012370

  • 匹配函数定义(vuln)
  • 匹配func call vuln(argv[1])

假如没有对vuln的调用，那就不回返回结果

image-20220803130323523

multi-query 实现

这块逻辑主要在 multi_query_worker ，即存在多个workitem的时候会触发，就是在匹配的时候会结合这些query，即将第一个query匹配到的结果先收集起来

1
2
3
4
5
6
7
8
9
let mut query_results = Vec::with_capacity(num_queries);
for _ in 0..num_queries {
    query_results.push(Vec::new());
}

// collect all results
for ctx in results_rx {
    query_results[ctx.query_index].push(ctx);
}

然后根据后面的query去做过滤，找到满足的pattern就打印出来

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
let filter = |x: &mut Vec<ResultsCtx>, y: &mut Vec<ResultsCtx>| {
        x.retain(|r| {
            y.iter()
                .any(|f| r.result.chainable(&r.source, &f.result, &f.source))
        })
    };

for i in 0..query_results.len() {
        let (part1, part2) = query_results.split_at_mut(i + 1);
        let a = part1.last_mut().unwrap();
        for b in part2 {
            filter(a, b);
            filter(b, a);
        }
    }

方便调试做的修改

1. 打印 query-tree 和 源码 AST方便定位问题

query-tree的话增加一个 -v 参数就行，会把query tree打印出来

少量代码测试这样是可以的，也可以使用log模块把信息打出来，不过数据太多了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
diff --git a/src/main.rs b/src/main.rs
index a819c5c..caf9a51 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -468,6 +468,8 @@ fn execute_queries_worker(
                 .enumerate()
                 .for_each(|(i, WorkItem { qt, identifiers: _ })| {
                     // Run query
+                    let tmp_tree = tree.root_node().to_sexp();
+                    info!("AST : {}", tmp_tree);
                     let matches = qt.matches(tree.root_node(), &source);
 
                     if matches.is_empty() {

直观多了 :

image-20220802153714210