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_list, domestic_dns_list, ipset_name="gfwlist"):
	"""
	生成 Dnsmasq 和 IPSet 规则文件。
	"""
	if not foreign_dns_list or not domestic_dns_list:
		log_info("错误：必须提供至少一个国内 DNS 和一个国外 DNS。")
		sys.exit(1)
		
	# --- 代理配置：通过本地 SOCKS5 隧道下载 GFWList ---
	proxies = {
		'http': f'socks5h://{proxy_server}',  # socks5h 使用远程 DNS 解析
		'https': f'socks5h://{proxy_server}'
	}
	log_info(f"配置代理：使用 {proxy_server} 通过隧道下载 GFWList...")
	# ----------------------------------------------------
		
	try:
		log_info(f"正在从 {GFWLIST_URL} 下载 GFWList...")
		# 传入 proxies 参数，强制 requests 走代理
		response = requests.get(GFWLIST_URL, timeout=30, proxies=proxies)
		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 失败，请检查 SSH SOCKS5 隧道状态！: {e}")
		# 如果下载失败，中止脚本以避免生成空文件
		sys.exit(1)

	domains = set()
	for line in content.splitlines():
		line = line.strip()
		if not line or line.startswith(('#', '!')):
			continue
		
		# 1. 清理前缀
		if line.startswith('||'):
			domain = line[2:]
		elif line.startswith('.'):
			domain = line[1:]
		else:
			domain = line
			
		# 2. 移除 URL 路径和查询参数
		if '/' in domain:
			domain = domain.split('/')[0]
			
		# 3. 过滤掉 IP 地址和通配符，只保留域名
		if domain and '.' in domain and not domain.startswith('*.'):
			# dnsmasq只接受不带前导点的域名，例如: example.com
			if domain.startswith('.'):
				domain = domain[1:]
				
			# 确保域名不包含非法字符
			if not any(c in domain for c in ('&', '?', '=', ':', '@', '%')):
				domains.add(domain)

	# 2. 生成 DNSMasq 配置文件
	dnsmasq_conf_path = "/etc/dnsmasq.d/gfwlist_router.conf"
	
	# 代理的上游 DNS (只使用提供的第一个国外 DNS)
	primary_foreign_dns = foreign_dns_list[0]
	proxy_port = proxy_server.split(':')[1]
	proxy_dns_server = "127.0.0.1#53053" 

	try:
		with open(dnsmasq_conf_path, 'w') as f:
			f.write(f"# Dnsmasq rules generated by {os.path.basename(__file__)} on {os.uname().nodename}\n")
			
			# 写入所有国内 DNS 服务器
			for dns in domestic_dns_list:
				f.write(f"server={dns}\n")
			
			# GFWList 域名设置
			for domain in sorted(list(domains)):
				# 1. 将域名加入 IPSet 集合 (由 dnsmasq 动态处理)
				f.write(f"ipset=/{domain}/{ipset_name}\n")
				# 2. 通过代理的 DNS 服务器解析 (流量会走 redsocks/ssh)
				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 集合存在 (必须在 Dnsmasq 启动前存在)
	try:
		# 使用 hash:ip 集合，如果集合已存在则清除其内容 (flush)
		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")
	
	# 允许接受多个参数 (nargs='+')
	parser.add_argument("-p", "--proxy", required=True, help="SOCKS5 Proxy server (e.g., 127.0.0.1:1086)")
	parser.add_argument("-f", "--foreign-dns", nargs='+', required=True, help="Foreign DNS IP(s) (e.g., 8.8.8.8 9.9.9.9)")
	parser.add_argument("-d", "--domestic-dns", nargs='+', required=True, help="Domestic DNS IP(s) (e.g., 223.5.5.5 114.114.114.114)")
	parser.add_argument("-n", "--name", default="gfwlist", help="IPSet Name")

	args = parser.parse_args()
	
	# 在开始生成规则之前，必须确保 requests 库已经安装 socks5 支持
	# 确保在 Bash 脚本中已经执行: sudo pip3 install requests PySocks
	
	log_info("开始生成 GFWList 规则...")
	generate_rules(args.proxy, args.foreign_dns, args.domestic_dns, args.name)
	log_info("规则生成完成。")