Fix loss of state, and a thorough test

This commit is contained in:
Neale Pickett 2020-09-24 20:23:43 -06:00
parent b9f024161f
commit 3a6c62985c
4 changed files with 32 additions and 24 deletions

View File

@ -2,13 +2,14 @@ package main
import ( import (
"fmt" "fmt"
"github.com/dirtbags/netshovel"
"github.com/google/gopacket"
"github.com/google/gopacket/tcpassembly"
"io" "io"
"log" "log"
"strings" "strings"
"sync" "sync"
"github.com/dirtbags/netshovel"
"github.com/google/gopacket"
"github.com/google/gopacket/tcpassembly"
) )
var wg sync.WaitGroup var wg sync.WaitGroup
@ -17,7 +18,7 @@ type SimpleStreamFactory struct {
} }
type SimpleStream struct { type SimpleStream struct {
netshovel.Stream *netshovel.Stream
} }
type SimplePacket struct { type SimplePacket struct {

View File

@ -31,23 +31,30 @@ func Shovel(factory tcpassembly.StreamFactory) {
assembler := tcpassembly.NewAssembler(streamPool) assembler := tcpassembly.NewAssembler(streamPool)
for _, fn := range flag.Args() { for _, fn := range flag.Args() {
handle, err := pcap.OpenOffline(fn) ShovelFile(fn, assembler)
if err != nil {
log.Fatal(err)
}
packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
packets := packetSource.Packets()
for packet := range packets {
if packet == nil {
break
}
if packet.NetworkLayer() == nil || packet.TransportLayer() == nil || packet.TransportLayer().LayerType() != layers.LayerTypeTCP {
continue
}
tcp := packet.TransportLayer().(*layers.TCP)
assembler.AssembleWithTimestamp(packet.NetworkLayer().NetworkFlow(), tcp, packet.Metadata().Timestamp)
}
} }
assembler.FlushAll() assembler.FlushAll()
} }
// ShovelFile shovels a single file.
// You must call assembler.FlushAll() at the end of this!
func ShovelFile(filename string, assembler *tcpassembly.Assembler) {
handle, err := pcap.OpenOffline(filename)
if err != nil {
log.Fatal(err)
}
packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
packets := packetSource.Packets()
for packet := range packets {
if packet == nil {
break
}
if packet.NetworkLayer() == nil || packet.TransportLayer() == nil || packet.TransportLayer().LayerType() != layers.LayerTypeTCP {
continue
}
tcp := packet.TransportLayer().(*layers.TCP)
assembler.AssembleWithTimestamp(packet.NetworkLayer().NetworkFlow(), tcp, packet.Metadata().Timestamp)
}
}

View File

@ -36,7 +36,7 @@ func TestHeaders(t *testing.T) {
t.Error("Uint8", fnord) t.Error("Uint8", fnord)
} }
biggun, err := pkt.Uint32LE("biggun") biggun, err := pkt.Uint32BE("biggun")
if err != nil { if err != nil {
t.Error(err) t.Error(err)
} }

View File

@ -43,8 +43,8 @@ type Stream struct {
// //
// You should embed Stream into your own Application protocol stream struct. // You should embed Stream into your own Application protocol stream struct.
// Use this to initialize the internal stuff netshovel needs. // Use this to initialize the internal stuff netshovel needs.
func NewStream(net, transport gopacket.Flow) Stream { func NewStream(net, transport gopacket.Flow) *Stream {
return Stream{ return &Stream{
Net: net, Net: net,
Transport: transport, Transport: transport,
conversation: make(chan Utterance, 100), conversation: make(chan Utterance, 100),