package main

import (
	"fmt"
	"log"

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

// PacketCapture 数据包捕获器
type PacketCapture struct {
	device      string
	handle      *pcap.Handle
	ruleManager *RuleManager
	logger      *Logger
	forwarder   *Forwarder
	stopChan    chan struct{}
}

// NewPacketCapture 创建新的数据包捕获器
func NewPacketCapture(device string, rm *RuleManager, logger *Logger, forwarder *Forwarder) *PacketCapture {
	return &PacketCapture{
		device:      device,
		ruleManager: rm,
		logger:      logger,
		forwarder:   forwarder,
		stopChan:    make(chan struct{}),
	}
}

// Start 启动数据包捕获
func (p *PacketCapture) Start() error {
	// 打开网络设备
	handle, err := pcap.OpenLive(p.device, 65536, true, pcap.BlockForever)
	if err != nil {
		return fmt.Errorf("无法打开网络设备: %v", err)
	}
	p.handle = handle

	// 设置BPF过滤器，捕获TCP、UDP和ICMP流量
	filter := "tcp or udp or icmp"
	if err := handle.SetBPFFilter(filter); err != nil {
		return fmt.Errorf("无法设置过滤器: %v", err)
	}

	// 开始捕获数据包
	go p.capturePackets()
	log.Printf("开始在设备 %s 上捕获数据包", p.device)
	return nil
}

// Stop 停止数据包捕获
func (p *PacketCapture) Stop() {
	close(p.stopChan)
	if p.handle != nil {
		p.handle.Close()
	}
	log.Println("停止数据包捕获")
}

// 捕获并处理数据包
func (p *PacketCapture) capturePackets() {
	packetSource := gopacket.NewPacketSource(p.handle, p.handle.LinkType())

	for {
		select {
		case packet := <-packetSource.Packets():
			p.processPacket(packet)
		case <-p.stopChan:
			return
		}
	}
}

// 处理捕获到的数据包
func (p *PacketCapture) processPacket(packet gopacket.Packet) {
	// 解析网络层和传输层
	ipLayer := packet.Layer(layers.LayerTypeIPv4)
	transportLayer := packet.Layer(layers.LayerTypeTCP)
	if transportLayer == nil {
		transportLayer = packet.Layer(layers.LayerTypeUDP)
	}
	icmpLayer := packet.Layer(layers.LayerTypeICMPv4)

	// 只处理包含IP层和传输层/ICMP层的数据包
	if ipLayer == nil || (transportLayer == nil && icmpLayer == nil) {
		return
	}

	// 提取IP信息
	ip, _ := ipLayer.(*layers.IPv4)
	srcIP := ip.SrcIP
	dstIP := ip.DstIP

	// 确定协议类型
	var protocol Protocol
	var srcPort, dstPort int

	if tcp, ok := transportLayer.(*layers.TCP); ok {
		protocol = ProtocolTCP
		srcPort = int(tcp.SrcPort)
		dstPort = int(tcp.DstPort)
	} else if udp, ok := transportLayer.(*layers.UDP); ok {
		protocol = ProtocolUDP
		srcPort = int(udp.SrcPort)
		dstPort = int(udp.DstPort)
	} else if icmpLayer != nil {
		protocol = ProtocolICMP
		srcPort = 0
		dstPort = 0
	} else {
		return
	}

	// 查找匹配的规则
	rule := p.ruleManager.MatchRule(srcIP, dstIP, srcPort, dstPort, protocol)

	// 应用规则
	if rule != nil {
		p.logger.LogPacket(rule, srcIP.String(), dstIP.String(), srcPort, dstPort, protocol, rule.Action)

		if rule.Action == ActionDeny {
			// 阻止数据包（不做任何处理）
			return
		}
	}

	// 如果允许通过且需要转发，则进行转发处理
	if rule == nil || rule.Action == ActionAllow {
		if p.forwarder.enabled {
			// 克隆数据包以便修改
			buf := gopacket.NewSerializeBuffer()
			err := gopacket.SerializePacket(buf, gopacket.SerializeOptions{FixLengths: true, ComputeChecksums: true}, packet)
			if err != nil {
				p.logger.Error("无法序列化数据包: ", err)
				return
			}

			// 解析克隆的数据包
			newPacket := gopacket.NewPacket(buf.Bytes(), packet.LinkLayer().LayerType(), gopacket.Default)
			newIpLayer := newPacket.Layer(layers.LayerTypeIPv4)
			newTransportLayer := newPacket.Layer(layers.LayerTypeTCP)
			if newTransportLayer == nil {
				newTransportLayer = newPacket.Layer(layers.LayerTypeUDP)
			}

			if newIpLayer != nil && newTransportLayer != nil {
					ip, _ := newIpLayer.(*layers.IPv4)
					transport, ok := newTransportLayer.(gopacket.TransportLayer)
					if !ok {
						p.logger.Error("Invalid transport layer type")
						return
					}
					p.forwarder.ForwardPacket(ip, transport, buf.Bytes())

				// 发送修改后的数据包
				if err := p.handle.WritePacketData(buf.Bytes()); err != nil {
					p.logger.Error("转发数据包失败: ", err)
				}
			}
		}
	}
}