使用Trie树快速搜索顶级域名

12 6月 2023
6 分钟阅读
标签:
trie

前言

在开发或者研究过程中经常会遇到域名，例如，攻击者可能会使用钓鱼攻击，欺骗用户点击看似合法但实际上是恶意的链接（最近zip域名比较火），这些链接中的域名都可能是伪造的。通过识别顶级域名，我们可以快速判断链接的真实性并减少恶意链接对系统的威胁。此外，在开发安全工具时，识别顶级域名也可以帮助我们更好地理解和分析网络流量，从而提高安全性。所以用rust基于trie树重写了https://github.com/john-kurkowski/tldextract/。

公共后缀列表

Mozilla维护了一个公共后缀列表，在页面中可以看到它在浏览器中的应用，还用各种编程语言的支持。得到公共后缀列表后初步对列表进行分析，可以得到列表以//开头作为注释，主要分为两类：ICANN，PRIVATE，以下面分割线分割。

// ===BEGIN ICANN DOMAINS===
...
// ===END ICANN DOMAINS===
...
// ===BEGIN PRIVATE DOMAINS===
...
// ===END PRIVATE DOMAINS===

每个后缀上面有注释标识来源和提交组织，例如：

// cn : https://en.wikipedia.org/wiki/.cn
// Submitted by registry [email protected]>
cn
ac.cn
com.cn
edu.cn
gov.cn
net.cn
org.cn
mil.cn
公司.cn
网络.cn
網絡.cn

还有一些特殊的后缀：
- 以*开头的为泛域名，就是前面加上任意一个词，拼接起来还是算一个顶级域名，比如：asd.kawasaki.jp。
- 以!开头的会被排除（为了约束上面的*），city.kawasaki.jp不算是顶级域名，www.city.kawasaki.jp提取出来的顶级域名是kawasaki.jp。

// jp geographic type names
// http://jprs.jp/doc/rule/saisoku-1.html
*.kawasaki.jp
*.kitakyushu.jp
*.kobe.jp
*.nagoya.jp
*.sapporo.jp
*.sendai.jp
*.yokohama.jp
!city.kawasaki.jp
!city.kitakyushu.jp
!city.kobe.jp
!city.nagoya.jp
!city.sapporo.jp
!city.sendai.jp
!city.yokohama.jp

Trie树是什么

Trie树是一种数据结构，用于高效地存储和搜索字符串集合。它通常用于在字符串集合中搜索前缀或匹配项。它是一种树形结构，其中每个节点代表一个字符。通过遍历树，可以构建完整的字符串。但是在域名这里明显用单个字符不太合适，因为域名是以.分割的，使用一个字符串作为最小单元构造树比较合理一点。
Trie树非常适合用于搜索前缀或匹配项，因为它可以在O(k)的时间内找到一个字符串，其中k是字符串的长度。trie树在许多应用程序中都有用，包括搜索引擎，拼写检查器和数据压缩。

Trie的结构

例如我将：下面四个域名转为树可以的到下面的树结构。

blog.kali-team.cn
www.gd.gov.cn
www.zj.gov.cn
mirrors.tuna.tsinghua.edu.cn

根节点没有任何词，每个节点有一个词，加粗的表示可以为顶级域名

flowchart TD
Root --> cn[<b>cn</b>]
Root --> ...
cn --> edu[<b>edu</b>]
cn --> gov[<b>gov</b>]
cn --> kali-team
kali-team --> blog
edu --> tsinghua
tsinghua --> tuna
tsinghua --> tsinghuawww[www]
tuna --> mirrors
gov --> gd
gov --> zj[zj]
gd --> www
zj --> zjwww[www]

将公共后缀列表构造到Trie树

了解了公共后缀列表和Trie树后，以cn为例子生成Trie树

cn
ac.cn
com.cn
edu.cn
gov.cn
net.cn
org.cn
mil.cn
公司.cn
网络.cn
網絡.cn

生成Trie树为：

flowchart TD
Root --> cn[<b>cn</b>]
cn --> edu[<b>edu</b>]
cn --> gov[<b>gov</b>]
cn --> com[<b>com</b>]
cn --> ac[<b>ac</b>]
cn --> net[<b>net</b>]
cn --> org[<b>org</b>]
cn --> mil[<b>mil</b>]
cn --> 公司[<b>公司</b>]
cn --> 网络[<b>网络</b>]
cn --> 網絡[<b>網絡</b>]

代码实现

实现一个嵌套的树结构，node为下一层的叶子节点，end表示当前这层是否可以为顶级域名。

/// TLDTrieTree
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TLDTrieTree {
    // 节点
    node: HashMap<String, TLDTrieTree>,
    // 是否可以为顶级域名
    end: bool,
}

插入数据

代码实现：

impl TLDTrieTree {
  /// Insert TLDTrieTree Construction Data
  #[inline]
  fn insert(&mut self, keys: Vec<&str>) {
    let keys_len = keys.len();
    let mut current_node = &mut self.node;
    for (index, mut key) in keys.clone().into_iter().enumerate() {
      let mut is_exclude = false;
      // 以!开头的需要排除掉
      if index == keys_len - 1 && key.starts_with('!') {
        key = &key[1..];
        is_exclude = true;
      }
      // 获取下一个节点，没有就插入默认节点
      let next_node = current_node.entry(key.to_string()).or_insert(TLDTrieTree {
        node: Default::default(),
        end: false,
      });
      // 当这是最后一个节点，设置可以为顶级域名
      if !is_exclude && (index == keys_len - 1)
                // 最后一个为*的，节点可以为顶级域名
                || (key != "*" && index == keys_len - 2 && keys[index + 1] == "*")
      {
        next_node.end = true;
      }
      current_node = &mut next_node.node;
    }
  }
}

以上面cn的edu.cn为例子插入Trie数，先将域名以.分割，再倒序得到列表[cn, edu]，按顺序全部插入得到。

{
  "node": {
    "cn": {
      "node": {
        "mil": {
          "node": {},
          "end": true
        },
        "com": {
          "node": {},
          "end": true
        },
        "xn--od0alg": {
          "node": {},
          "end": true
        },
        "xn--io0a7i": {
          "node": {},
          "end": true
        },
        "gov": {
          "node": {},
          "end": true
        },
        "xn--55qx5d": {
          "node": {},
          "end": true
        },
        "net": {
          "node": {},
          "end": true
        },
        "ac": {
          "node": {},
          "end": true
        },
        "edu": {
          "node": {},
          "end": true
        },
        "org": {
          "node": {},
          "end": true
        }
      },
      "end": true
    }
  },
  "end": false
}

将下面后缀构造为Trie数后得到

*.ck
!www.ck

ck同级的end为true，表示可以为顶级域名，但是www前面有!，所以end为false。

{
  "node": {
    "ck": {
      "node": {
        "*": {
          "node": {},
          "end": true
        },
        "www": {
          "node": {},
          "end": false
        }
      },
      "end": true
    }
  },
  "end": false
}

查询数据

代码实现：

impl TLDTrieTree {
/// Search tree, return the maximum path searched
  #[inline]
  fn search(&self, keys: &[String]) -> Vec<Suffix> {
    let mut suffix_list = Vec::new();
    let mut current_node = &self.node;
    for key in keys.iter() {
      match current_node.get(key) {
        Some(next_node) => {
          suffix_list.push(Suffix {
            suffix: key.to_string(),
            end: next_node.end,
          });
          current_node = &next_node.node;
        }
        None => {
          if let Some(next_node) = current_node.get("*") {
            suffix_list.push(Suffix {
              suffix: key.to_string(),
              end: next_node.end,
            });
          }
          break;
        }
      }
    }
    suffix_list
  }
}

查询时将域名以.分割后倒序逐级从右边往左边搜索，不管end是否为true，直到找不到最后一个词，进入match中的None分支，再判断节点是否为存在*，如果存在将返回的节点的end设置为true，得到一个带有节点和是否可以为顶级域名属性的列表。
例如：www.asd.city.kawasaki.jp以.分割后倒序得到[jp, kawasaki, city, asd, www]，直到搜索到city，找不到下一层了，返回后缀列表：

[Suffix { suffix: "jp", end: true }, Suffix { suffix: "kawasaki", end: true }, Suffix { suffix: "city", end: false }]

然后将这个列表pop数据，遇到第一个end为true的Suffix，当前Suffix和剩下的组成顶级域名：也就是：kawasaki.jp

ExtractResult {
        subdomain: Some(
            "www.asd",
        ),
        domain: Some(
            "city",
        ),
        suffix: Some(
            "kawasaki.jp",
        ),
        registered_domain: Some(
            "city.kawasaki.jp",
        ),
    },

结论

在本文中，我们介绍了如何使用Trie树来快速查找顶级域名。这种方法可以帮助我们快速识别恶意链接，并提高我们的安全性。我们还介绍了Trie树的基本概念和实现细节。希望这篇文章对你有所帮助！

使用Trie树快速搜索顶级域名

目录

前言

公共后缀列表

Trie树是什么

Trie的结构

将公共后缀列表构造到Trie树

代码实现

插入数据

查询数据

结论

参考