IDA 中的大规模路径搜索方法

背景

本文主要解决的是这么一个问题: 在 IDA 中如何查找两个函数之间的调用路径？这问题似乎很简单，IDA 工具本身就提供了查看调用链路的功能，懂得 IDAPython 的也可以很容易写出下面的程序：

def find_path(start, end, path=None):
    path = path or []
    if start == end:
        return path
    path.append(end)
    for xref in idautils.CodeRefsTo(end, False):
            caller = ida_funcs.get_func(xref)
            return find_path(start, caller.start_ea, path)
1
2
3
4
5
6
7
8

不过，这样就行了吗？

缺陷

上述代码在大部分情况下基本都能实现目标，但其中存在几个缺陷。

首先，其中没有考虑函数本身循环/递归调用的情况，在遇到调用环时会进入死循环。这个问题好解决，只需要在全局定义一个 set，然后将每次遇到的节点（函数）存放进去，并且每次遇到新节点时都会先判断一下当前节点是否已经在 set 之中，如果在则表示当前路径存在环，可以考虑提前退出。

这个问题解决了，还有另外一个问题。上述代码本质上是通过递归调用 IDA 的交叉引用接口去查找上一级路径的，这在代码不多的时候问题不大，可一旦遇到过长的调用链路，就可能出现递归调用过深导致栈耗尽而出现异常。

更为雪上加霜的是，使用递归会使得我们实际的搜索算法是深度优先的，因此即便有很短的调用链路，可能也会因为节点遍历顺序靠后而无法搜索到。

双栈算法

为了解决递归搜索引起的栈溢出问题，就需要将搜索方法切换为非递归的算法。读者可能已经意识到了，寻找调用路径的问题，其实可以抽象为图论中的寻路问题。更准确地说，是有向图中的寻路问题。

如果是查找最短路径，我们可以用 Dijkstra 算法或者 A-星 算法。不过笔者最初的目标是查找所有路径，因此用到了另一个算法，称为双栈算法。其大致流程如下：

1. 建立两个栈，起始节点放到主栈，起始节点的相邻结点（以列表形式）放到辅栈。使得主栈和辅栈的高度相等；
2. 从辅栈栈顶的列表中取出一个元素（节点），并压到主栈中。此时主栈比辅栈多一层，因此辅栈需要放入主栈栈顶的相邻节点。
3. 此时如果辅栈栈顶中的节点有已经在主栈中的，需要将其拿掉；
4. goto 2，直到辅栈栈顶为空列表，将主栈和辅栈的空列表同时移除，这个过程称为削栈；
5. goto 2，继续重复建栈和削栈过程，直到主栈的栈顶节点是目标节点，此时表示找到了一条路径，即主栈中栈底到栈顶的所有元素。
6. 继续重复上述过程，直到主栈为空，即找到了所有的路径。

关于该算法的具体实现原理可以参考这篇文章。

该算法的优势是通过循环和双栈结构替换递归实现查找，不用担心后者空间耗尽的问题。同时在第 3 步的时候由于在辅栈中去掉了与主栈栈顶重复的结点，也巧妙地避免带环的路径。

实现

Talk is cheap？下面是笔者实现的完整代码：

import ida_funcs
import idautils

class Finder(object):
    def __init__(self):
        self.s1 = [] # main stack
        self.s2 = [] # neighbor stack
        self.visited = {}

    def get_neighbors(self, node):
        neighbors = set()
        for xref in idautils.CodeRefsTo(node, False):
            caller = ida_funcs.get_func(xref)
            neighbors.add(caller.start_ea)
        return neighbors

    def build_dual(self, node):
        self.s1.append(node)
        self.visited[node] = True
        neighbors = self.get_neighbors(node)
        for n in list(neighbors):
            if self.visited.get(n) is True:
                neighbors.remove(n)
        self.s2.append(neighbors)
        # assert(len(self.s1) == len(self.s2))

    def cut_dual(self):
        self.s2.pop()
        node = self.s1.pop()
        self.visited[node] = False

    def find(self, source, sink):
        print("start searching path {:#x} -> {:#x}".format(source, sink))
        self.s1.clear()
        self.s2.clear()
        # build dual stack
        self.build_dual(sink)
        while len(self.s1) > 0:
            neighbors = self.s2.pop()
            if len(neighbors) > 0:
                new_node = neighbors.pop()
                self.s2.append(neighbors) # remains
                self.build_dual(new_node)
            else:
                # empty neighbors
                self.s2.append(None) # dummy
                self.cut_dual()
                continue
            if self.s1[-1] == source:
                # found a path
                yield list(reversed(self.s1))
                self.cut_dual()
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
39
40
41
42
43
44
45
46
47
48
49
50
51
52

使用时只需要指定起始函数地址和末尾（目标）函数地址即可。例如，打印二者中的所有路径：

e = Finder()
for path in e.find(start, end):
    print(path)
1
2
3

案例分析

下面看几个路径搜索的具体案例。二进制文件应该是来自个 CTF 比赛题目，但是具体题目链接不记得了。

faster1

IDA 打开，main 函数打开长这样：

void sub_42AE87()
{
  if ( counter++ )
    exit(-1);
  switch ( read_int() )
  {
    case 0:
      sub_403272();
      break;
    case 1:
      sub_422274();
      break;
    case 2:
      sub_433C45();
      break;
    case 3:
      sub_417D30();
      break;
    case 4:
      sub_41191B();
      break;
    case 5:
      sub_4206A2();
      break;
    case 6:
      sub_42CCEB();
      break;
    case 7:
      sub_419C84();
      break;
    case 8:
      sub_41B326();
      break;
    case 9:
      sub_431443();
      break;
    default:
      sub_42AE87();
      break;
  }
}
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
39
40
41

read_int 是用户输入一个数字，其中每个 case 中的函数都长得差不多，如下：

void sub_403272()
{
  int v0; // eax

  v0 = counter++;
  if ( v0 != 343 )
    exit(-1);
  switch ( read_int() )
  {
    case 0:
      sub_43189C();
      break;
    case 1:
      sub_411F32();
      break;
    // ....
    case 9:
      sub_422DC7();
      break;
    default:
      sub_42AE87();
      break;
  }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

每个函数开头会判断 counter 是否匹配，匹配才进入下一个。最后一个函数称为 win，其中存在一个栈溢出可以出发漏洞，漏洞利用部分不是重点，这里先忽略。该函数调用者为：

void sub_40CACF()
{
  int v0; // eax

  v0 = counter++;
  if ( v0 != 999 )
    exit(-1);
  switch ( read_int() )
  {
    // ...
    case 9:
      win();
      break;
    default:
      sub_42AE87();
      break;
  }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

所以题目的逻辑很简单，只需要从 main 函数调用到上述函数时候计数正好是 999 即可，由于每层加一，也就是需要调用层数正好为 1000，让我们反推出输入的数字组合。

这道题我最初是通过编写 IDAPython 脚本不断反向寻址找出的路径解出的，因为有限制条件，因此不会出现路径爆炸的问题。使用 Finder 也可以寻找出路径，只需要增加前级的约束，即过滤相邻节点即可。

faster2

这题是在上一题的基础上做了一点修改，主要改动是每个函数开头不再检测 counter 的值，而只是单纯加一，每个函数基本都长这样：

void sub_40FD44()
{
  ++counter;
  switch ( read_int() )
  {
    case 0:
      sub_41F786();
      break;
    case 1:
      sub_40BF4E();
      break;
    //...9
    default:
      sub_4098AE();
      break;
  }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

但是，最后在 win 函数中检查了 counter 的值：

ssize_t win()
{
  char buf[207]; // [rsp+0h] [rbp-D0h] BYREF
  char v2; // [rsp+CFh] [rbp-1h]

  if ( counter != 1000 )
    exit(-1);
  v2 = getchar();
  if ( v2 != 10 )
    exit(-1);
  write(1, "WOW,U R GREAT !\n", 0x10uLL);
  return read(0, buf, 0x200uLL);
}
1
2
3
4
5
6
7
8
9
10
11
12
13

也就是说，要求从起始到末尾的调用层数同样是 1000，由于每个函数少了 counter 的判断和约束，因此回溯时就无法判断和过滤无效分支，只能通过暴力去搜索。前文中使用的递归搜索方法在遇到这种量级的层数调用时候毫无疑问会耗尽栈空间而失败。

值得一提的是，在使用 Finder 进行搜索时，因为时间关系无法直接找到层数正好的调用链路，但可以找到许多有效路径。由于在程序中存在大量的环形调用，因此我们可以随便找到一个大小合适的环，只需要其长度倍数与有效路径之和满足条件即可形成一条长度为 1000 的调用链路。具体解题细节就不在此赘述了。

小结

本文主要是记录和分享了一种在 IDA 中通过非递归去实现的路径搜索算法，其算法核心是将递归的搜索替换为栈+循环的方式，可以应用在大规模的程序中避免递归内存耗尽。另外通过修改 get_neighbors 方法也可以方便地拓展到 Ghidra 或者 BinaryNinja 中。其实这也不算什么很新颖的东西，不过能将算法应用到自己写的脚本中感觉还是很奇妙的。