#!/usr/bin/env python3
# standalone_smart_dns.py - 轻量级 DNS 转发服务器，具有防污染能力

import socket
import threading
import ipaddress
import os
import time
import requests
import json
from dnslib import DNSRecord, DNSHeader, RR, QTYPE, A, AAAA

# 配置项
LISTEN_IP = "127.0.0.1"
LISTEN_PORT = 53
CHINA_DNS = "114.114.114.114"
FOREIGN_DNS = "1.1.1.1"
CHINA_IP_LIST_URL = "https://raw.githubusercontent.com/mayaxcn/china-ip-list/master/chnroute.txt"
CHINA_IPV6_LIST_URL = "https://raw.githubusercontent.com/mayaxcn/china-ip-list/master/chnroute_v6.txt"
CHINA_IP_CACHE_FILE = "/tmp/.china_ip_list.txt"
CHINA_IPV6_CACHE_FILE = "/tmp/.china_ipv6_list.txt"
CACHE_DURATION = 7 * 24 * 60 * 60  # 7 天秒数
MAX_PACKET_SIZE = 1024
TIMEOUT = 5

# 全局变量，存储 IP 网络
cn_networks = []
cn_networks_v6 = []

def get_china_ip_list(is_ipv6=False):
    """获取或更新中国 IP 列表"""
    url = CHINA_IPV6_LIST_URL if is_ipv6 else CHINA_IP_LIST_URL
    cache_file = CHINA_IPV6_CACHE_FILE if is_ipv6 else CHINA_IP_CACHE_FILE
    
    # 检查缓存是否存在且未过期
    if os.path.exists(cache_file):
        file_age = time.time() - os.path.getmtime(cache_file)
        if file_age < CACHE_DURATION:
            with open(cache_file, 'r') as f:
                return [line.strip() for line in f if line.strip() and not line.startswith('#')]
    
    # 下载新列表
    try:
        response = requests.get(url)
        if response.status_code == 200:
            ip_list = [line.strip() for line in response.text.split('\n') 
                      if line.strip() and not line.startswith('#')]
            
            # 保存到缓存
            with open(cache_file, 'w') as f:
                f.write('\n'.join(ip_list))
            
            return ip_list
        else:
            # 如果下载失败但缓存存在，使用缓存
            if os.path.exists(cache_file):
                with open(cache_file, 'r') as f:
                    return [line.strip() for line in f if line.strip() and not line.startswith('#')]
            else:
                return []
    except Exception as e:
        print(f"获取中国 IP 列表出错: {e}")
        if os.path.exists(cache_file):
            with open(cache_file, 'r') as f:
                return [line.strip() for line in f if line.strip() and not line.startswith('#')]
        return []

def is_cn_ip(ip):
    """检查 IP 是否在中国网络列表中"""
    global cn_networks, cn_networks_v6
    
    try:
        ip_obj = ipaddress.ip_address(ip)
        
        if ip_obj.version == 4:
            # 首次运行时加载中国 IPv4 列表
            if not cn_networks:
                china_ip_list = get_china_ip_list(is_ipv6=False)
                cn_networks = [ipaddress.ip_network(net) for net in china_ip_list]
            
            # 检查 IP 是否在任何网络中
            for net in cn_networks:
                if ip_obj in net:
                    return True
        elif ip_obj.version == 6:
            # 首次运行时加载中国 IPv6 列表
            if not cn_networks_v6:
                china_ipv6_list = get_china_ip_list(is_ipv6=True)
                cn_networks_v6 = [ipaddress.ip_network(net) for net in china_ipv6_list]
            
            # 检查 IP 是否在任何网络中
            for net in cn_networks_v6:
                if ip_obj in net:
                    return True
        
        return False
    except:
        return False

def forward_dns_request(domain, dns_server, record_type):
    """将 DNS 请求转发到指定的 DNS 服务器"""
    try:
        # 正确使用 dnslib 创建 DNS 请求
        from dnslib import DNSQuestion
        record = DNSRecord()
        record.add_question(DNSQuestion(domain, getattr(QTYPE, record_type)))
        packet = record.pack()
        
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.settimeout(TIMEOUT)
        
        s.sendto(packet, (dns_server, 53))
        reply_packet, _ = s.recvfrom(MAX_PACKET_SIZE)
        
        reply = DNSRecord.parse(reply_packet)
        return reply
    except Exception as e:
        print(f"查询 {dns_server} 关于 {domain} ({record_type}) 出错: {e}")
        return None

def handle_dns_request(data, client_addr, sock):
    """处理 DNS 请求并返回智能选择的响应"""
    try:
        request = DNSRecord.parse(data)
        qname = str(request.q.qname)
        qtype = QTYPE[request.q.qtype]
        
        print(f"查询: {qname} ({qtype}) 来自 {client_addr[0]}")
        
        # 仅对 A 和 AAAA 记录使用智能逻辑
        if qtype in ('A', 'AAAA'):
            # 查询两个 DNS 服务器
            cn_reply = forward_dns_request(qname, CHINA_DNS, qtype)
            foreign_reply = forward_dns_request(qname, FOREIGN_DNS, qtype)
            
            # 提取回复中的 IP
            cn_ips = []
            if cn_reply and cn_reply.rr:
                cn_ips = [str(rr.rdata) for rr in cn_reply.rr if rr.rtype == getattr(QTYPE, qtype)]
            
            foreign_ips = []
            if foreign_reply and foreign_reply.rr:
                foreign_ips = [str(rr.rdata) for rr in foreign_reply.rr if rr.rtype == getattr(QTYPE, qtype)]
            
            # 检查任何 CN IP 是否实际在中国
            cn_ips_in_china = [ip for ip in cn_ips if is_cn_ip(ip)]
            
            # 智能选择逻辑
            if cn_ips_in_china:
                # 中国 DNS 返回了中国 IP，使用 CN 结果
                print(f"为 {qname} ({qtype}) 使用中国 DNS")
                response = cn_reply
            elif foreign_ips:
                # 中国 DNS 没有返回中国 IP，但国外 DNS 有结果
                # 可能是 DNS 污染，使用国外结果
                print(f"为 {qname} ({qtype}) 使用国外 DNS")
                response = foreign_reply
            elif cn_ips:
                # 默认使用中国 DNS 结果（如果有）
                print(f"默认为 {qname} ({qtype}) 使用中国 DNS")
                response = cn_reply
            else:
                # 如果以上都失败，使用国外 DNS 或空响应
                print(f"默认为 {qname} ({qtype}) 使用国外 DNS")
                response = foreign_reply if foreign_reply else cn_reply
        else:
            # 对于非 A/AAAA 记录，直接转发到国外 DNS
            response = forward_dns_request(qname, FOREIGN_DNS, qtype)
            if not response:
                response = forward_dns_request(qname, CHINA_DNS, qtype)
        
        # 如果有响应，发送它
        if response:
            # 保持原始请求 ID
            response.header.id = request.header.id
            sock.sendto(response.pack(), client_addr)
        else:
            # 创建空响应
            response = DNSRecord(DNSHeader(id=request.header.id, qr=1, aa=0, ra=1), q=request.q)
            sock.sendto(response.pack(), client_addr)
            
    except Exception as e:
        print(f"处理请求出错: {e}")

def run_dns_server():
    """运行 DNS 服务器"""
    # 尝试 UDP 和 TCP 套接字
    try:
        # 用于 DNS 的 UDP 套接字
        udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp_sock.bind((LISTEN_IP, LISTEN_PORT))
        print(f"DNS 服务器监听 {LISTEN_IP}:{LISTEN_PORT} (UDP)")
        
        while True:
            data, addr = udp_sock.recvfrom(MAX_PACKET_SIZE)
            threading.Thread(target=handle_dns_request, args=(data, addr, udp_sock)).start()
            
    except Exception as e:
        print(f"服务器错误: {e}")
    finally:
        if 'udp_sock' in locals():
            udp_sock.close()

if __name__ == "__main__":
    # 加载库
    try:
        import dnslib
    except ImportError:
        print("安装必要的包...")
        import pip
        pip.main(['install', 'dnslib', 'requests'])
        import dnslib
    
    print("启动智能 DNS 服务器...")
    run_dns_server()