import argparse
import base64
import requests
import os
import sys

# GFWList 官方（或常见）的地址
GFWLIST_URL = "https://raw.githubusercontent.com/gfwlist/gfwlist/master/gfwlist.txt"

def log_info(message):
    print(f"[INFO] {message}")

def generate_rules(proxy_server, foreign_dns, domestic_dns, ipset_name="gfwlist"):
    """
    生成 Dnsmasq 和 IPSet 规则文件。
    """
    try:
        log_info(f"正在从 {GFWLIST_URL} 下载 GFWList...")
        response = requests.get(GFWLIST_URL, timeout=30)
        response.raise_for_status()
        
        # 1. 解码 GFWList
        try:
            content = base64.b64decode(response.text).decode('utf-8')
        except:
            # 如果内容不是 base64 编码，则直接使用
            content = response.text
            
    except requests.exceptions.RequestException as e:
        log_info(f"下载 GFWList 失败，使用本地缓存：{e}")
        # 允许继续运行，依赖旧的规则文件
        content = ""

    domains = set()
    for line in content.splitlines():
        line = line.strip()
        if not line or line.startswith(('#', '!')):
            continue
        
        # 提取域名（忽略 IP、通配符规则和特殊标记）
        if line.startswith('||'):
            domain = line[2:]
        elif line.startswith('.'):
            domain = line[1:]
        elif '/' in line:
            continue
        else:
            domain = line

        # 简单过滤，确保是域名
        if domain and '.' in domain and not domain.startswith('*.'):
            domains.add(domain)

    # 2. 生成 DNSMasq 配置文件
    dnsmasq_conf_path = "/etc/dnsmasq.d/gfwlist_router.conf"
    
    # 确保 foreign_dns 包含端口 (例如 127.0.0.1#1086)
    proxy_dns_server = f"{foreign_dns}#{proxy_server.split(':')[1]}" 

    try:
        with open(dnsmasq_conf_path, 'w') as f:
            f.write(f"# Dnsmasq rules generated by gfw_rules_generator.py\n")
            
            # 设置国内 DNS
            f.write(f"server={domestic_dns}\n")
            
            # GFWList 域名设置
            for domain in sorted(list(domains)):
                # 1. 将域名加入 IPSet 集合
                f.write(f"ipset=/{domain}/{ipset_name}\n")
                # 2. 通过代理的 DNS 服务器解析
                f.write(f"server=/{domain}/{proxy_dns_server}\n")
                
        log_info(f"已生成 Dnsmasq 规则到 {dnsmasq_conf_path}，包含 {len(domains)} 个域名。")
    except IOError as e:
        log_info(f"写入 Dnsmasq 文件失败: {e}")
        sys.exit(1)

    # 3. 确保 IPSet 存在
    # 脚本不负责添加 IP 地址（由 Dnsmasq 在运行时动态添加）
    # 但需要确保 IPSet 集合本身存在。
    try:
        os.system(f"sudo ipset create {ipset_name} hash:ip family inet hashsize 1024 maxelem 65536 &>/dev/null || sudo ipset flush {ipset_name}")
        log_info(f"已确保 IPSet 集合 '{ipset_name}' 存在。")
    except Exception as e:
        log_info(f"创建 IPSet 失败: {e}")
        sys.exit(1)
    
    return True

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="GFWList Rules Generator for Dnsmasq/IPSet")
    parser.add_argument("-p", "--proxy", required=True, help="SOCKS5 Proxy server (e.g., 127.0.0.1:1086)")
    parser.add_argument("-f", "--foreign-dns", required=True, help="Foreign DNS IP (e.g., 8.8.8.8)")
    parser.add_argument("-d", "--domestic-dns", required=True, help="Domestic DNS IP (e.g., 223.5.5.5)")
    parser.add_argument("-n", "--name", default="gfwlist", help="IPSet Name")

    args = parser.parse_args()
    
    log_info("开始生成 GFWList 规则...")
    generate_rules(args.proxy, args.foreign_dns, args.domestic_dns, args.name)
    log_info("规则生成完成。")