第35天：DNS深入 - 递归与迭代查询

# day35_dns_query_tool.py
# 功能：DNS查询工具

import socket
import struct
from typing import List, Tuple, Dict
import time


class DNSQuery:
    """
    DNS查询工具

    功能：
        1. 构造DNS查询报文
        2. 发送查询
        3. 解析响应
        4. 支持多种记录类型
    """

    # DNS记录类型
    QTYPE = {
        'A': 1,      # IPv4地址
        'NS': 2,     # 域名服务器
        'CNAME': 5,  # 规范名称
        'SOA': 6,    # 授权开始
        'PTR': 12,   # 指针记录
        'MX': 15,    # 邮件交换
        'TXT': 16,   # 文本
        'AAAA': 28,  # IPv6地址
        'ANY': 255   # 所有记录
    }

    def __init__(self, dns_server='8.8.8.8', port=53):
        """
        初始化DNS查询器

        参数：
            dns_server: DNS服务器地址
            port: DNS端口（默认53）
        """
        self.dns_server = dns_server
        self.port = port

    def build_query(self, domain: str, qtype: str = 'A') -> bytes:
        """
        构造DNS查询报文

        参数：
            domain: 要查询的域名
            qtype: 查询类型（A, AAAA, MX等）

        返回：
            DNS查询报文（字节）

        DNS报文格式：
        +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
        |                    Header                     |
        +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
        |                   Question                    |
        +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
        |                    Answer                     |
        +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
        |                  Authority                    |
        +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
        |                  Additional                   |
        +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
        """
        # 1. 构造Header（12字节）
        transaction_id = 0x1234  # 事务ID
        flags = 0x0100  # 标准查询，期望递归
        questions = 1   # 问题数
        answers = 0     # 回答数
        authority = 0   # 权威记录数
        additional = 0  # 附加记录数

        header = struct.pack(
            '!HHHHHH',
            transaction_id, flags,
            questions, answers, authority, additional
        )

        # 2. 构造Question
        # 域名编码：将 www.example.com 编码为 \x03www\x07example\x03com\x00
        qname = b''
        for part in domain.split('.'):
            qname += bytes([len(part)]) + part.encode()
        qname += b'\x00'  # 结束标记

        qtype_value = self.QTYPE.get(qtype.upper(), 1)
        qclass = 1  # IN (Internet)

        question = qname + struct.pack('!HH', qtype_value, qclass)

        return header + question

    def parse_response(self, response: bytes) -> Dict:
        """
        解析DNS响应

        参数：
            response: DNS响应报文

        返回：
            解析后的结果字典
        """
        result = {
            'header': {},
            'questions': [],
            'answers': [],
            'authority': [],
            'additional': []
        }

        # 解析Header
        header = struct.unpack('!HHHHHH', response[:12])
        result['header'] = {
            'id': header[0],
            'flags': header[1],
            'questions': header[2],
            'answers': header[3],
            'authority': header[4],
            'additional': header[5]
        }

        # 解析Answer部分（简化版本）
        offset = 12

        # 跳过Question部分
        while response[offset] != 0:
            offset += 1
        offset += 5  # 跳过结束符和QTYPE、QCLASS

        # 解析Answer
        for _ in range(result['header']['answers']):
            # 简化：直接提取IP地址（假设是A记录）
            if offset + 12 <= len(response):
                # 跳过名称（指针）
                offset += 2

                # 读取类型、类、TTL、数据长度
                rtype, rclass, ttl, rdlength = struct.unpack(
                    '!HHIH',
                    response[offset:offset+10]
                )
                offset += 10

                # 读取数据
                rdata = response[offset:offset+rdlength]
                offset += rdlength

                if rtype == 1 and rdlength == 4:  # A记录
                    ip = '.'.join(str(b) for b in rdata)
                    result['answers'].append({
                        'type': 'A',
                        'ttl': ttl,
                        'data': ip
                    })

        return result

    def query(self, domain: str, qtype: str = 'A', verbose: bool = True) -> Dict:
        """
        执行DNS查询

        参数：
            domain: 域名
            qtype: 查询类型
            verbose: 是否显示详细信息

        返回：
            查询结果
        """
        if verbose:
            print(f"\n正在查询: {domain} ({qtype}记录)")
            print(f"DNS服务器: {self.dns_server}:{self.port}")

        # 构造查询
        query_data = self.build_query(domain, qtype)

        # 创建UDP socket
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.settimeout(3)

        try:
            # 发送查询
            start_time = time.time()
            sock.sendto(query_data, (self.dns_server, self.port))

            # 接收响应
            response, _ = sock.recvfrom(512)
            query_time = (time.time() - start_time) * 1000

            # 解析响应
            result = self.parse_response(response)
            result['query_time'] = query_time

            if verbose:
                print(f"查询时间: {query_time:.2f}ms")
                print(f"回答数量: {result['header']['answers']}")

                if result['answers']:
                    print("\n结果:")
                    for answer in result['answers']:
                        print(f"  类型: {answer['type']}")
                        print(f"  TTL: {answer['ttl']}秒")
                        print(f"  数据: {answer['data']}")

            return result

        except socket.timeout:
            if verbose:
                print("❌ 查询超时")
            return {'error': 'timeout'}
        except Exception as e:
            if verbose:
                print(f"❌ 查询失败: {e}")
            return {'error': str(e)}
        finally:
            sock.close()

    def query_with_system(self, domain: str, verbose: bool = True):
        """
        使用系统DNS查询（对比）

        参数：
            domain: 域名
            verbose: 是否显示详细信息
        """
        if verbose:
            print(f"\n使用系统DNS查询: {domain}")

        try:
            start_time = time.time()
            result = socket.getaddrinfo(domain, None)
            query_time = (time.time() - start_time) * 1000

            if verbose:
                print(f"查询时间: {query_time:.2f}ms")
                print("\n结果:")
                for item in result:
                    if item[0] == socket.AF_INET:  # IPv4
                        print(f"  IPv4: {item[4][0]}")
                    elif item[0] == socket.AF_INET6:  # IPv6
                        print(f"  IPv6: {item[4][0]}")

            return {'results': result, 'query_time': query_time}

        except Exception as e:
            if verbose:
                print(f"❌ 查询失败: {e}")
            return {'error': str(e)}


def compare_dns_servers():
    """对比不同DNS服务器的性能"""
    print("\n" + "="*60)
    print(" DNS服务器性能对比")
    print("="*60)

    dns_servers = [
        ('Google DNS', '8.8.8.8'),
        ('Cloudflare DNS', '1.1.1.1'),
        ('114 DNS', '114.114.114.114'),
    ]

    test_domains = [
        'www.baidu.com',
        'www.google.com',
        'www.github.com'
    ]

    print(f"\n测试域名: {', '.join(test_domains)}\n")

    for name, server in dns_servers:
        print(f"\n{name} ({server}):")
        print("-" * 60)

        query_tool = DNSQuery(server)
        total_time = 0
        success_count = 0

        for domain in test_domains:
            result = query_tool.query(domain, verbose=False)
            if 'query_time' in result:
                total_time += result['query_time']
                success_count += 1
                print(f"  {domain}: {result['query_time']:.2f}ms")

        if success_count > 0:
            avg_time = total_time / success_count
            print(f"\n  平均查询时间: {avg_time:.2f}ms")


def main():
    """主程序"""
    print("""
    ╔══════════════════════════════════════════════════════════╗
    ║       DNS查询工具 v1.0                                    ║
    ║       DNS Query Tool                                     ║
    ╚══════════════════════════════════════════════════════════╝
    """)

    # 创建DNS查询器
    dns_query = DNSQuery(dns_server='8.8.8.8')

    # 演示1：基本查询
    print("\n【演示1：基本DNS查询】")
    dns_query.query('www.baidu.com', 'A')

    input("\n按Enter继续...")

    # 演示2：系统查询对比
    print("\n【演示2：自定义vs系统DNS查询对比】")
    dns_query.query_with_system('www.baidu.com')

    input("\n按Enter继续...")

    # 演示3：DNS服务器性能对比
    print("\n【演示3：不同DNS服务器性能对比】")
    compare_dns_servers()

    print("""
\n总结：
DNS查询过程包括：
1. 构造查询报文
2. 发送UDP查询（端口53）
3. 接收响应
4. 解析结果

优化建议：
- 使用快速DNS服务器
- 合理设置TTL
- 启用DNS缓存
- 考虑使用DNS over HTTPS (DoH)

明天我们将学习DNS安全（DNSSEC）！
    """)


if __name__ == "__main__":
    main()