Python 多线程 DNS 搜索性能优化

时间：2024-04-03 18:25:42

Python中的多线程经常用于IO密集型任务，如网络请求，其中DNS查询是常见的一种场景。由于全局解释器锁（GIL）的存在，Python的多线程并不适合计算密集型任务，但对于IO密集型任务，如DNS查询，多线程可以显著提高性能。那么如果遇到下面的问题，可以通过这样的解决方法解决。

在这里插入图片描述

1、问题背景

原有 Python DNS 搜索代码在扫描大范围 IP 地址时速度较慢，需要进行优化以提高性能。同时，使用多线程会导致写入文件时出现问题，需要找到一种方法来解决这个问题。

2、解决方案

  1. 优化 DNS 查询过程：
  • 优化 DNS 查询包的生成和发送过程，减少不必要的操作。
  • 调整超时时间以减少等待时间。
  1. 优化多线程处理：
  • 使用线程池来管理线程，提高线程利用率。
  • 使用锁来控制对文件写入的访问，避免多线程写入冲突。
  1. 使用异步 I/O：
  • 将文件写入操作改为异步 I/O，以提高 I/O 性能。
  1. 代码示例：
import socket
import struct
import threading
import os
import sys
import time
import asyncio

# 基本 DNS 头部结构，用于 1 个查询
def build_dns_query(host):
    packet = struct.pack("!HHHHHH", 0x0001, 0x0100, 1, 0, 0, 0)

    for name in host:
        query = struct.pack("!b" + str(len(name)) + "s", len(name), name)
        packet = packet + query

    packet = packet + struct.pack("!bHH", 0, 1, 1)

    return packet


# 测试查询，用于 www.google.com
TEST_QUERY = build_dns_query(["www", "google", "com"])
DNS_PORT = 53
TIMEOUT = 2


# 扫描服务器的 DNS
async def scan_dns(addr, timeout):
    reader, writer = await asyncio.open_connection(addr, DNS_PORT)

    # 发送 DNS 查询请求
    send_count = writer.write(TEST_QUERY)
    if send_count <= 0:
        return False

    # 等待响应
    try:
        data = await reader.read(1024)
    except asyncio.TimeoutError:
        return False

    return True


# 将 IP 地址解析为整型元组
def extract_ip(ip):
    partip = ip.split(".")
    if len(partip) != 4:
        print("Invalid IP address:", ip)

    try:
        ip_tuple = (int(partip[0]), int(partip[1]), int(partip[2]), int(partip[3]))
    except ValueError:
        print("Invalid IP address:", ip)

    return ip_tuple


# 主函数
async def main(start_ip, end_ip):
    # 存储找到的 DNS 服务器
    found_dns = []

    # 扫描所有 IP 地址
    for i0 in range(start_ip[0], end_ip[0] + 1):
        for i1 in range(start_ip[1], end_ip[1] + 1):
            for i2 in range(start_ip[2], end_ip[2] + 1):
                for i3 in range(start_ip[3], end_ip[3] + 1):
                    # 构建 IP 地址
                    ip_addr = f"{i0}.{i1}.{i2}.{i3}"

                    print(f"Scanning {ip_addr}...", end=" ")

                    # 扫描地址
                    ret = await scan_dns(ip_addr, TIMEOUT)

                    if ret:
                        found_dns.append(ip_addr)
                        print("Found!")
                        await write_file(ip_addr)
                    else:
                        print("")

    print("Found DNS servers:", found_dns)


# 写入文件
async def write_file(ip_addr):
    file = open("dns_servers.txt", "a")
    file.write(ip_addr + "\n")
    file.close()


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python dnsfind.py <start_ip> <end_ip>")
        exit()

    # 转换 IP 地址到整型元组
    start_ip = extract_ip(sys.argv[1])
    end_ip = extract_ip(sys.argv[2])

    # 执行主函数
    asyncio.run(main(start_ip, end_ip))

根据你的应用和机器的具体情况调整线程池的大小。对于高并发的DNS查询，使用异步IO（如asyncio库）可能比多线程更有效率。例如dnspython提供的异步解析功能，可能比使用socket.gethostbyname更高效。实现这些优化策略后，你应该能够显著提高Python程序中DNS查询的性能。如果有更好的建议欢迎评论区留言讨论。

相关文章