test_tcpassembly.go

package main

import (
	"bufio"
	"log"
	"os"

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

func NewAssembler() *tcpassembly.Assembler {
	streamFactory := &myFactory{}
	streamPool := tcpassembly.NewStreamPool(streamFactory)
	return tcpassembly.NewAssembler(streamPool)
}

type myFactory struct {
	counter int64
}

func (f *myFactory) New(net, transport gopacket.Flow) tcpassembly.Stream {
	f.counter += 1
	s := &myStream{i: f.counter}
	log.Printf("stream #%d: (created)", s.i)
	return s
}

type myStream struct {
	i int64
}

func (s *myStream) Reassembled(reass []tcpassembly.Reassembly) {
	log.Printf("stream #%d: Reassembled", s.i)
}
func (s *myStream) ReassemblyComplete() {
	log.Printf("stream #%d: ReassemblyComplete", s.i)
}

func main() {
	file, err := os.Open(os.Args[1])
	if err != nil {
		log.Fatal(err)
	}

	pcap, err := pcapgo.NewReader(bufio.NewReader(file))
	if err != nil {
		log.Fatal(err)
	}

	packetSource := gopacket.NewPacketSource(pcap, pcap.LinkType())
	packets := packetSource.Packets()

	assembler := NewAssembler()

	var packetNo int64
	packetNo = 1
	for packet := range packets {
		if packet.NetworkLayer() == nil || packet.TransportLayer() == nil || packet.TransportLayer().LayerType() != layers.LayerTypeTCP {
			continue
		}
		tcp := packet.TransportLayer().(*layers.TCP)
		log.Printf("packet %d: calling AssembleWithTimestamp()", packetNo)
		packetNo++
		assembler.AssembleWithTimestamp(packet.NetworkLayer().NetworkFlow(), tcp, packet.Metadata().Timestamp)
	}

	log.Print("calling FlushAll")

	assembler.FlushAll()
}