package pcap

import (
	"errors"
	"fmt"
	"github.com/gopacket/gopacket"
	"github.com/gopacket/gopacket/ip4defrag"
	"github.com/gopacket/gopacket/layers"
	"github.com/gopacket/gopacket/pcapgo"
	"io"
	"os"
	"testing"

	"gotest.tools/v3/assert"
)

func Test(t *testing.T) {
	file, err := os.Open("fragmented.pcapng")
	pcapr, err := pcapgo.NewNgReader(file, pcapgo.DefaultNgReaderOptions)
	assert.NilError(t, err)

	defragmenter := ip4defrag.NewIPv4Defragmenter()

	for {
		data, ci, err := pcapr.ZeroCopyReadPacketData()
		if errors.Is(err, io.EOF) {
			fmt.Println("End of file")
			return
		}
		assert.NilError(t, err)

		packet := gopacket.NewPacket(data, pcapr.LinkType(), gopacket.Default)
		if ipv4Layer := packet.Layer(layers.LayerTypeIPv4); ipv4Layer != nil {
			ipv4LayerCasted := ipv4Layer.(*layers.IPv4)
			defragmented, err := defragmenter.DefragIPv4WithTimestamp(ipv4LayerCasted, ci.Timestamp)
			assert.NilError(t, err)
			if defragmented == nil || defragmented.Protocol != layers.IPProtocolUDP {
				continue // packet fragment, we don't have whole packet yet.
			}

			udpPacket := gopacket.NewPacket(defragmented.LayerPayload(), layers.LayerTypeUDP, gopacket.Default)
			if udpLayer := udpPacket.Layer(layers.LayerTypeUDP); udpLayer != nil {
				payload := udpLayer.LayerPayload()
				assert.Equal(t, len(payload) > 0, true)
				fmt.Println("Read payload success")
			}
		}
	}
}