package pcaphelper

import (
	"io"
	"net"
	"time"

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

// OnTrHook 发射包的钩子函数
type OnTrHook func([]byte) []byte

// WriteRaw 将给定的数据注入 pcap 句柄
func (p *PcapHelper) WritePacketData(bytes []byte) error {
	return p.handle.WritePacketData(bytes)
}

// PlayPCAPFile 播放 pcap 文件，支持 bpf 过滤文件中的报文，可以设置或重置发包间隔，通过 hook 修改每一个数据包的内容
// 当 interval<0 将无间隔播放；当 interval=0 按 pcap 文件中默认的间隔播放，当 interval>0 时按 interval 间隔播放
func (p *PcapHelper) PlayPCAPFileWithHook(file, bpfExpr string, interval time.Duration, onTr OnTrHook) error {
	handleRead, err := pcap.OpenOffline(file)
	if err != nil {
		return errors.Wrap(err, "open PCAP file")
	}

	if err := handleRead.SetBPFFilter(bpfExpr); err != nil {
		return errors.Wrap(err, "set BPF filter on handle reader")
	}

	var (
		sentCount        int
		lastPktTimestamp time.Time
	)
	for {
		// read packet from handle
		data, ci, err := handleRead.ReadPacketData()
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return errors.Wrap(err, "failed to read packet")
		}
		if ci.CaptureLength != ci.Length {
			continue // do not write truncated packets
		}

		// interval time
		if sentCount != 0 {
			switch {
			case interval == 0:
				time.Sleep(ci.Timestamp.Sub(lastPktTimestamp))
			case interval > 0:
				time.Sleep(interval)
			case interval < 0:
				// no pause
			}
		}

		data = onTr(data)
		if data == nil {
			continue
		}

		if err := p.WritePacketData(data); err != nil {
			return errors.Wrap(err, "failed to write packet")
		}
		lastPktTimestamp = ci.Timestamp
		sentCount++
	}
}

// PlayPCAPFile 播放 pcap 文件，支持 bpf 过滤文件中的报文，可以设置或重置发包间隔
func (p *PcapHelper) PlayPCAPFile(file, bpfExpr string, interval time.Duration) error {
	return p.PlayPCAPFileWithHook(file, bpfExpr, interval, func(bytes []byte) []byte { return bytes })
}

// MakeHookMutateEth 修改以太网层的 来源MAC 和 目的MAC
func MakeTrHookThatMutateEth(srcMAC, dstMAC net.HardwareAddr) OnTrHook {
	return func(bytes []byte) []byte {
		var (
			pkt        = gopacket.NewPacket(bytes, layers.LayerTypeEthernet, gopacket.Default)
			eth        = new(layers.Ethernet)
			ethPayload = make(gopacket.Payload, 0, len(bytes))
			buffer     = gopacket.NewSerializeBuffer()
		)

		if ethLayer := pkt.Layer(layers.LayerTypeEthernet); ethLayer != nil {
			eth = ethLayer.(*layers.Ethernet)

			if srcMAC != nil {
				eth.SrcMAC = srcMAC
			}
			if dstMAC != nil {
				eth.DstMAC = dstMAC
			}

			if eth.Payload != nil {
				ethPayload = eth.Payload
			}

			_ = gopacket.SerializeLayers(buffer, gopacket.SerializeOptions{ComputeChecksums: true, FixLengths: true},
				eth, ethPayload,
			)
			return buffer.Bytes()
		}

		return bytes
	}
}

// MakeHookMutateEthIP4 修改以太网层的 来源MAC 和 目的MAC，IP4 层的 源IP 和 目的IP
func MakeTrHookThatMutateEthIP4(srcMAC, dstMAC net.HardwareAddr, srcIP, dstIP net.IP) OnTrHook {
	return func(bytes []byte) []byte {
		var (
			pkt        = gopacket.NewPacket(bytes, layers.LayerTypeEthernet, gopacket.Default)
			eth        = new(layers.Ethernet)
			ip4        = new(layers.IPv4)
			ip4Payload = make(gopacket.Payload, 0, len(bytes))
			buffer     = gopacket.NewSerializeBuffer()
		)

		if ethLayer := pkt.Layer(layers.LayerTypeEthernet); ethLayer != nil {
			eth = ethLayer.(*layers.Ethernet)

			if srcMAC != nil {
				eth.SrcMAC = srcMAC
			}
			if dstMAC != nil {
				eth.DstMAC = dstMAC
			}
		}

		if ip4Layer := pkt.Layer(layers.LayerTypeIPv4); ip4Layer != nil {
			ip4 = ip4Layer.(*layers.IPv4)

			if srcIP != nil {
				ip4.SrcIP = srcIP
			}
			if dstIP != nil {
				ip4.DstIP = dstIP
			}

			if ip4.Payload != nil {
				ip4Payload = ip4.Payload
			}

			_ = gopacket.SerializeLayers(buffer, gopacket.SerializeOptions{ComputeChecksums: true, FixLengths: true},
				eth, ip4, ip4Payload,
			)
			return buffer.Bytes()
		}

		return bytes
	}
}