package main

import (
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/google/gopacket/pcapgo"
)

// 打印可用的网络接口信息 (保持不变)
func printAvailableInterfaces() {
	devices, err := pcap.FindAllDevs()
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println("可用的网络接口:")
	for _, device := range devices {
		fmt.Printf("名称: %s\n描述: %s\n地址: %+v\n\n", device.Name, device.Description, device.Addresses)
	}
}

// ----------------------------------------------------------------------
// 移除 isIPInList 函数 (不再需要)
// ----------------------------------------------------------------------

// 处理捕获到的数据包 (原 printPacketInfo 的优化版)
func handlePacket(packet gopacket.Packet) {
	// 1. 快速提取网络层 (IPv4)
	// 使用 Layer() 比 Layers() 稍微快一点点，但如果追求极致性能，建议使用 ZeroCopy 模式
	ipLayer := packet.Layer(layers.LayerTypeIPv4)
	if ipLayer == nil {
		return
	}
	ip, ok := ipLayer.(*layers.IPv4)
	if !ok {
		return
	}

	// 2. 提取源 IP 字符串
	srcIP := ip.SrcIP.String()

	// 3. 将 IP 推入全局处理队列
	// 这里不再加锁，也不判断是否存在链表中，直接交给 PushIPToQueue 的非阻塞逻辑处理
	// PushIPToQueue 会负责去重 (PendingIPs) 和限流
	PushIPToQueue(srcIP)

	// 注意：为了性能，去掉了这里的日志打印。
	// 在高流量攻击下，每秒打印几千行日志会导致 IO 阻塞，拖慢抓包速度。
	// 具体的封禁日志会在 processIP (消费者) 中打印。
}

// startPacketCapture 启动数据包捕获
func startPacketCapture() {
	if *InterfaceName == "" {
		log.Fatal("未指定网络接口，请使用 -i 参数")
	}

	// 打开指定的网络接口进行实时数据包捕获
	// snaplen 设置为 1600 足够捕获头部信息，没必要设为 65535，可以稍微节省内存
	handle, err := pcap.OpenLive(*InterfaceName, 1600, true, pcap.BlockForever)
	if err != nil {
		log.Fatalf("打开网络接口 %s 出错: %v", *InterfaceName, err)
	}
	defer func() {
		fmt.Println("清理抓包资源...")
		handle.Close()
	}()

	// 设置 BPF 过滤器
	if *Protocol != "" {
		if err = handle.SetBPFFilter(*Protocol); err != nil {
			log.Fatalf("设置 BPF 过滤器出错: %v", err)
		}
	}

	// --- PCAP 文件保存逻辑 (如果不需要存包，这部分其实是最大的性能瓶颈) ---
	var pcapWriter *pcapgo.Writer
	if *PcapFile != "" {
		file, err := os.Create(*PcapFile)
		if err != nil {
			log.Fatalf("创建输出文件出错: %v", err)
		}
		defer file.Close()

		pcapWriter = pcapgo.NewWriter(file)
		// 写入文件头
		if err = pcapWriter.WriteFileHeader(1600, layers.LinkTypeEthernet); err != nil {
			log.Fatalf("写入全局头部出错: %v", err)
		}
	}

	// 使用 ZeroCopyPacketDataSource 可以减少内存拷贝，提升性能 (Linux Only)
	// 如果在非 Linux 环境编译报错，请改回 gopacket.NewPacketSource
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	// packetSource.NoCopy = true // 开启 NoCopy 模式 (可选，稍微危险但更快)

	log.Printf(" 正在监听网络接口 %s, 使用过滤器 '%s'...\n", *InterfaceName, *Protocol)

	// 创建信号通道
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

	// 启动 Goroutine 处理数据包
	go func() {
		// 直接遍历 channel
		for packet := range packetSource.Packets() {

			// 1. 核心业务：提取 IP 进队列
			handlePacket(packet)

			// 2. 可选业务：存盘
			// 警告：在高并发攻击下，写文件 IO 会成为瓶颈。建议生产环境非必要不开启 -o 参数
			if pcapWriter != nil {
				// 注意：CaptureInfo 和 Data 必须在 packet 被复用前读取
				err := pcapWriter.WritePacket(packet.Metadata().CaptureInfo, packet.Data())
				if err != nil {
					log.Printf("写入数据包出错: %v", err)
				}
			}
		}
	}()

	// 阻塞等待信号
	<-sigChan
	fmt.Println("\n停止抓包...")
}