package main

import (
	"container/list"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"
)

var BuildDate = "unknown" // 由编译时注入

func init() {
	// 强制使用 Go 的纯用户态 DNS 解析器
	net.DefaultResolver = &net.Resolver{
		PreferGo: true,
	}
}

var (
	daemon = flag.Bool("d", false, "守护进程模式")
	child  = flag.Bool("child", false, "子进程模式, (不要使用！！！)")
)

// 全局变量
var (
	InterfacesList bool    // 是否列出网络接口
	InterfaceName  *string // 网络接口名称
	PcapFile       *string // 输出文件名
	Protocol       *string // BPF 过滤器

	IPSET_NUMBER   int    // 当前使用的 ipset 集合编号
	MAX_IPSET_NAME = 100  // 最大 ipset 集合数量
	IPSET_NAME     string // 当前使用的 ipset 集合名称

	IpList  = list.New() // 存储 IPv4 地址的链表
	IpMutex sync.Mutex   // 保护 ipList 的互斥锁

	ProcessedIPMap = map[string]time.Time{} // 使用 map 存储已处理的 IP
	ProcessedMutex sync.Mutex               // 互斥锁保护 ProcessedIPMap

	local_ipv4_addr string
)

// 启动子进程
func StartChildProcess() (*exec.Cmd, error) {
	args := []string{}
	for _, arg := range os.Args[1:] {
		if !strings.HasPrefix(arg, "-child") { // 只过滤 -child 标志
			args = append(args, arg)
		}
	}
	args = append(args, "-child=true")
	cmd := exec.Command(os.Args[0], args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("启动子进程失败: %w", err)
	}
	log.Printf(" 子进程已启动, PID: %d\n", cmd.Process.Pid)
	return cmd, nil
}

// 停止子进程
func StopChildProcess(cmd *exec.Cmd) error {
	if cmd == nil || cmd.Process == nil {
		return fmt.Errorf("子进程无效")
	}

	// 尝试优雅地停止子进程
	if err := cmd.Process.Signal(syscall.SIGTERM); err != nil {
		return fmt.Errorf("发送 SIGTERM 信号失败: %w", err)
	}

	done := make(chan error, 1)
	go func() {
		done <- cmd.Wait()
	}()

	select {
	case err := <-done:
		if err != nil {
			return fmt.Errorf("等待子进程退出时出错: %w", err)
		}
	case <-time.After(5 * time.Second): // 超时时间调整为5秒
		fmt.Println("子进程未在规定时间内退出，尝试强制终止...")
		if err := cmd.Process.Kill(); err != nil {
			return fmt.Errorf("强制终止子进程失败: %w", err)
		}
	}

	fmt.Printf("子进程已停止, PID: %d\n", cmd.Process.Pid)
	return nil
}

// 等待信号并优雅退出
func waitForSignalAndCleanUp(cmd *exec.Cmd) {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

	for sig := range sigChan {
		fmt.Printf("主进程收到信号: %v\n", sig)
		if cmd != nil && cmd.Process != nil {
			_ = cmd.Process.Signal(sig) // 转发信号到子进程
		}
		if sig == syscall.SIGINT || sig == syscall.SIGTERM {
			break
		}
	}

	if err := StopChildProcess(cmd); err != nil {
		log.Fatalf("清理子进程时遇到错误: %v", err)
	}
	fmt.Println("主进程退出")
}

// 守护进程模式
func StartDaemon() {
	// 创建一个新的实例并让当前进程退出。注意这并不是守护进程的标准实现。
	args := []string{}
	for _, arg := range os.Args[1:] {
		if !strings.HasPrefix(arg, "-d") && !strings.HasPrefix(arg, "-child") { // 过滤掉 -d 和 -child 标志
			args = append(args, arg)
		}
	}
	args = append(args, "-d=false", "-child=false")

	cmd := exec.Command(os.Args[0], args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		log.Fatalf("无法启动新实例: %v", err)
	}
	fmt.Println("新实例已启动，当前进程将退出")
	os.Exit(0)
}

// 子进程逻辑
func RunChildProcess() {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

	ticker := time.NewTicker(3 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case sig := <-sigChan:
			fmt.Printf("子进程收到信号: %v, 准备退出...\n", sig)
			return
		case <-ticker.C:

		}
	}
}

// 保存ProcessedIPMap到文件
func saveMapToFile(filePath string) error {
	ProcessedMutex.Lock()
	defer ProcessedMutex.Unlock()

	file, err := os.Create(filePath)
	if err != nil {
		return fmt.Errorf("创建文件失败: %w", err)
	}
	defer file.Close()

	encoder := json.NewEncoder(file)
	if err := encoder.Encode(ProcessedIPMap); err != nil {
		return fmt.Errorf("编码 ProcessedIPMap 失败: %w", err)
	}

	//log.Println("ProcessedIPMap 已成功保存到文件")
	return nil
}

func loadFromFile(filePath string, logMessage string) error {
	ProcessedMutex.Lock()
	defer ProcessedMutex.Unlock()

	file, err := os.Open(filePath)
	if err != nil {
		if os.IsNotExist(err) {
			log.Println("文件不存在，跳过加载 Map")
			return nil
		}
		return fmt.Errorf("打开文件失败: %w", err)
	}
	defer file.Close()

	// 尝试用新格式解码
	decoder := json.NewDecoder(file)
	temp := make(map[string]time.Time)
	if err := decoder.Decode(&temp); err == nil {
		ProcessedIPMap = temp
		log.Println(logMessage)
		return nil
	}

	// 如果失败，尝试旧格式
	file.Seek(0, 0) // 重置读取位置
	decoder = json.NewDecoder(file)
	oldTemp := make(map[string]string)
	if err := decoder.Decode(&oldTemp); err == nil {
		for ip := range oldTemp {
			//ProcessedIPMap[ip] = time.Now() // 给旧 IP 打个当前时间戳
			ProcessedIPMap[ip] = time.Date(1971, 1, 1, 0, 0, 0, 0, time.UTC) // 标记为永不过期
		}
		log.Println(logMessage + "（从旧格式转换）")
		return nil
	}

	return fmt.Errorf("解码 Map 失败: %w", err)
}

func InitMap() {
	if err := loadFromFile("cn.json", " Map 已成功从文件加载"); err != nil {
		log.Fatalf("  加载 Map 失败: %v", err)
	}

	defer func() {
		// 程序退出时保存数据
		if err := saveMapToFile("cn.json"); err != nil {
			log.Printf("  保存 Map 失败: %v", err)
		}
	}()

}

func WriteLocalAddr() {
	local_ipv4_addr = GetLocalIpv4Addr() // 本机地址
	// 将本机外网地址加入到已处理集合中
	if local_ipv4_addr != "NULL" {
		//log.Printf("\033[33m %s 本机地址 \033[0m\n", ipStr)
		ProcessedMutex.Lock()
		//ProcessedIPMap[local_ipv4_addr] = time.Now()
		ProcessedIPMap[local_ipv4_addr] = time.Date(1971, 1, 1, 0, 0, 0, 0, time.UTC) // 标记为永不过期
		ProcessedMutex.Unlock()
	}

	// 写入json文件
	if err := saveMapToFile("cn.json"); err != nil {
		log.Printf("实时保存 Map 失败: %v", err)
	}

}

func RunMainProcess() { // 主进程逻辑

	log.Println(" 主进程启动...")

	WriteLocalAddr() // 将本机外网地址加入到已处理集合中

	cmd, err := StartChildProcess()
	if err != nil {
		log.Fatalf("子进程启动失败: %v", err)
	}

	IPSET_NUMBER = 0
	IPSET_NAME = fmt.Sprintf("root%d", IPSET_NUMBER)
	if return_value := Is_Name_Ipset(IPSET_NAME); return_value == 1 {
		createIPSet(IPSET_NAME)
	}

	// 启动抓包
	go startPacketCapture()

	// 启动IP地域判断，管理
	go func() {
		for {

			IpMutex.Lock()        // 锁定互斥锁
			if IpList.Len() > 0 { // 链表不为空
				e1 := IpList.Front() // 获取链表第一个元素
				ipStr := e1.Value.(net.IP).String()
				region, _ := ip2region(ipStr) // 离线库初步判断地域

				ProcessedMutex.Lock()
				_, processed := ProcessedIPMap[ipStr] // 检查是否已处理
				ProcessedMutex.Unlock()

				if processed { // 如果尚未处理
					log.Printf("\033[33m %s 已经标记为国内，跳过！！！ \033[0m\n", ipStr)
					IpList.Remove(e1)
					goto next
				}

				// 检查是否在白名单中
				if _, ip_ := whiteList[ipStr]; ip_ {
					log.Printf("\033[33m %s 跳过白名单, 跳过！！！ \033[0m\n", ipStr)
					IpList.Remove(e1)

					setName, err := RemoveIPIfInSets("root", MAX_IPSET_NAME, ipStr)
					if err != nil {
						log.Printf(" %s 删除 IP 出错: %v\n", ipStr, err)
					} else if setName != "" {
						log.Printf(" %s 已从 %s 中移除 \n", ipStr, setName)
					} else {
						log.Printf(" %s 不在任何 IPSet 中，无需移除\n", ipStr)
					}

					goto next
				}

				REGION := "中国 内网" // 默认地域

				if Is_Ip_Ipset(ipStr) != 0 { // IP 不在 ipset 集合中

					//if !strings.Contains(region, "中国") && !strings.Contains(region, "内网") {				// 离线库判断不在中国内
					if !ContainsPart(region, REGION) {
						log.Printf("\033[33m [%s %s] 离线库为国外, 进一步API判断\033[0m\n", ipStr, region)

						if position, err := curl_(ipStr); err != nil { // 使用 API 判断地域
							log.Printf("获取IP地域出错: %v", err)
						} else {
							log.Printf("\033[31m [%s %s]\033[0m\n\n", ipStr, position) // 打印地域

							//if !strings.Contains(position, "中国") && !strings.Contains(position, "内网") {	// API 判断为国外
							if !ContainsPart(region, REGION) {
								AddIPSet(IPSET_NAME, ipStr) // 添加 IP 到 ipset 集合

								// 钉钉告警，废弃！钉钉可能限制文本长度，和发送次数！
								// warning_ding(ipStr, position)    // 警告 IP 地域
							} else {
								log.Printf("\033[33m %s 离线库为国外, API 判断为国内, 标记为已处理\033[0m\n", ipStr)

								// 将 IP 标记为已处理，国内地址
								ProcessedMutex.Lock()
								ProcessedIPMap[ipStr] = time.Now()
								ProcessedMutex.Unlock()

								// 写入json文件
								if err := saveMapToFile("cn.json"); err != nil {
									log.Printf(" 实时保存 Map 失败: %v", err)
								}
							}
						}
					}

				} else { // IP 已在 ipset 集合中
					log.Printf("\033[31m %s 在 ipset 集合中 \033[0m\n", ipStr)
				}

				// 无论是否已处理，都移除该 IP
				IpList.Remove(e1)
			}
		next:
			IpMutex.Unlock() // 解锁互斥锁

			// 打印当前info
			log.Printf(" 当前Ip链表长度:%d, Ipset名:%s, 长:%d, ProcessedIPMap:[%s]当前长度:%d\n", IpList.Len(), IPSET_NAME, func() int { // 打印 当前 Ipset 链长度
				_len, _err := NumIPSet(IPSET_NAME)
				if _err == nil {
					return _len
				}
				return 0
			}(), ProcessedIPMap, len(ProcessedIPMap))

			time.Sleep(1 * time.Second) // 防止高频运行
		}
	}()

	// 启动一个 goroutine（后台任务），用于定期清理 ProcessedIPMap 中过期的 IP 条目
	go func() {
		for {
			time.Sleep(1 * time.Minute)

			// 获取当前时间，用于与 map 中的时间戳进行比较
			now := time.Now()

			// 加锁，确保并发安全地访问全局变量 ProcessedIPMap
			ProcessedMutex.Lock()

			// 遍历 ProcessedIPMap，检查每个 IP 的记录是否已超过 24 小时
			for ip, t := range ProcessedIPMap {

				if t.Year() == 1971 {
					continue // 不清理标记为“永不过期”的 IP
				}

				// 如果当前时间减去记录时间大于 1 小时，则删除该条目
				if now.Sub(t) > 30*time.Minute {
					delete(ProcessedIPMap, ip)
				}
			}

			// 解锁，允许其他 goroutine 访问 ProcessedIPMap
			ProcessedMutex.Unlock()

			// 打印日志，表示本次清理已完成
			log.Println(" 已清理过期 ProcessedIPMap 项")
		}
	}()

	// 定时重新加载白名单
	go func() {
		for {
			time.Sleep(1 * time.Minute) // 每 10 分钟自动刷新
			if err := LoadWhiteList("whitelist.txt"); err != nil {
				log.Printf(" 刷新白名单失败: %v", err)
			}
		}
	}()

	// 启动防火墙管理
	go func() {
		for {
			if ipset_len, _ := NumIPSet(IPSET_NAME); ipset_len >= 65535 {
				log.Printf("\033[31m ipset %s 列表已满 %d \033[0m\n", IPSET_NAME, ipset_len)

				// 创建新的 ipset 集合
				IPSET_NUMBER++

				if IPSET_NUMBER >= MAX_IPSET_NAME {
					//log.Printf("已创建 %d 个集合!!!", MAX_IPSET_NAME)
					log.Printf("\033[31m 已创建 %d 个集合!!! \033[0m\n", MAX_IPSET_NAME)
				}

				IPSET_NAME = fmt.Sprintf("root%d", IPSET_NUMBER)
				if return_value := Is_Name_Ipset(IPSET_NAME); return_value == 1 {
					createIPSet(IPSET_NAME)
				}

			}

			time.Sleep(3 * time.Second)
		}
	}()

	// 等待信号并清理
	waitForSignalAndCleanUp(cmd)
}

func HandleCmd() {
	// 定义命令行标志
	var instruction string
	var help bool
	InterfaceName = flag.String("i", "", "指定要使用的网络接口")
	flag.BoolVar(&InterfacesList, "l", false, "列出可用的网络接口")
	Protocol = flag.String("f", "tcp", "指定 BPF 过滤器")
	PcapFile = flag.String("o", "", "保存捕获数据的输出文件（可选）")
	flag.StringVar(&instruction, "s", "", "-s start 启动 Iptables 规则\n-s stop  停止 Iptables 规则\n-s list  打印 Iptables 规则")
	flag.BoolVar(&help, "h", false, "")
	flag.BoolVar(&help, "help", false, "帮助信息")
	flag.Parse()

	if help {
		fmt.Printf(
			"\t\tDenyip firewall\n"+
				"\tVersion 0.2\n"+
				"\tE-mail: aixiao@aixiao.me\n"+
				"\tBuild Date: %s\n", BuildDate)

		flag.Usage()
		fmt.Printf("\n")
		os.Exit(0)
	}

	if instruction != "" {
		switch instruction {
		case "start":
			for i := 0; i < MAX_IPSET_NAME; i++ {
				_name := fmt.Sprintf("root%d", i)
				iptables_add(_name)
			}

			os.Exit(0)
		case "stop":
			for i := 0; i < MAX_IPSET_NAME; i++ {
				_name := fmt.Sprintf("root%d", i)
				iptables_del(_name)
			}
			os.Exit(0)
		case "l":
			fallthrough
		case "list":
			_ = iptables_list()
			os.Exit(0)
		default:
			log.Fatalf("未知的操作: %s. 请使用 'start' 或 'stop'.", instruction)
		}
	}

	if InterfacesList {
		printAvailableInterfaces()
		os.Exit(0)
	}

	if *InterfaceName == "" {
		log.Fatal("请使用 -i 标志指定网络接口，或者使用 -l 列出接口。")
	}
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU()) // 设置最大CPU核数
	HandleCmd()

	embed_ip2region()

	// 加载白名单
	err := LoadWhiteList("whitelist.txt")
	if err != nil {
		log.Fatalf(" whiteList Map 加载白名单失败: %v", err)
	} else {
		log.Println(" whiteList Map 白名单加载成功")
	}

	// 守护进程模式
	if *daemon {
		StartDaemon()
	}

	// 子进程逻辑
	if *child {
		RunChildProcess() // 子进程逻辑
		return
	}

	InitMap()
	RunMainProcess() // 主进程逻辑

}
