added payload flag to preserve payloads for TCP, UDP and USB packets

dreadl0ck · Jan 13, 2019 · 655a29f · 655a29f
1 parent 5d7c91a
commit 655a29f
Show file tree

Hide file tree

Showing 13 changed files with 924 additions and 838 deletions.
diff --git a/cmd/flags.go b/cmd/flags.go
@@ -58,4 +58,5 @@ var (
 
 	flagBaseLayer     = flag.String("base", "ethernet", "select base layer")
 	flagDecodeOptions = flag.String("opts", "lazy", "select decoding options")
+	flagPayload       = flag.Bool("payload", false, "capture payload for supported layers")
 )
diff --git a/cmd/main.go b/cmd/main.go
@@ -150,6 +150,7 @@ func main() {
 		return
 	}
 
+	encoder.CapturePayload = *flagPayload
 	printLogo()
 
 	// print configuration as table

diff --git a/encoder/shared.go b/encoder/shared.go
@@ -27,6 +27,9 @@ var (
 	// BlockSize is the file system block size
 	BlockSize int
 
+	// CapturePayload for encoders that support it
+	CapturePayload = false
+
 	allEncoderNames = make(map[string]struct{})
 	errorMap        *AtomicCounterMap
 )

diff --git a/encoder/tcp.go b/encoder/tcp.go
@@ -22,7 +22,13 @@ import (
 
 var tcpEncoder = CreateLayerEncoder(types.Type_NC_TCP, layers.LayerTypeTCP, func(layer gopacket.Layer, timestamp string) proto.Message {
 	if tcp, ok := layer.(*layers.TCP); ok {
-		var opts []*types.TCPOption
+		var (
+			opts    []*types.TCPOption
+			payload []byte
+		)
+		if CapturePayload {
+			payload = layer.LayerPayload()
+		}
 		for _, o := range tcp.Options {
 			opts = append(opts, &types.TCPOption{
 				OptionData:   o.OptionData,
@@ -53,6 +59,7 @@ var tcpEncoder = CreateLayerEncoder(types.Type_NC_TCP, layers.LayerTypeTCP, func
 			Options:        opts,
 			PayloadEntropy: Entropy(tcp.Payload),
 			PayloadSize:    int32(len(tcp.Payload)),
+			Payload:        payload,
 		}
 	}
 	return nil

diff --git a/encoder/udp.go b/encoder/udp.go
@@ -22,6 +22,10 @@ import (
 
 var udpEncoder = CreateLayerEncoder(types.Type_NC_UDP, layers.LayerTypeUDP, func(layer gopacket.Layer, timestamp string) proto.Message {
 	if udp, ok := layer.(*layers.UDP); ok {
+		var payload []byte
+		if CapturePayload {
+			payload = layer.LayerPayload()
+		}
 		return &types.UDP{
 			Timestamp:      timestamp,
 			SrcPort:        int32(udp.SrcPort),
@@ -30,6 +34,7 @@ var udpEncoder = CreateLayerEncoder(types.Type_NC_UDP, layers.LayerTypeUDP, func
 			Checksum:       int32(udp.Checksum),
 			PayloadEntropy: Entropy(udp.Payload),
 			PayloadSize:    int32(len(udp.Payload)),
+			Payload:        payload,
 		}
 	}
 	return nil

diff --git a/encoder/usb.go b/encoder/usb.go
@@ -22,6 +22,10 @@ import (
 
 var usbEncoder = CreateLayerEncoder(types.Type_NC_USB, layers.LayerTypeUSB, func(layer gopacket.Layer, timestamp string) proto.Message {
 	if usb, ok := layer.(*layers.USB); ok {
+		var payload []byte
+		if CapturePayload {
+			payload = layer.LayerPayload()
+		}
 		return &types.USB{
 			Timestamp:              timestamp,
 			ID:                     uint64(usb.ID),
@@ -42,6 +46,7 @@ var usbEncoder = CreateLayerEncoder(types.Type_NC_USB, layers.LayerTypeUSB, func
 			UrbStartFrame:          uint32(usb.UrbStartFrame),
 			UrbCopyOfTransferFlags: uint32(usb.UrbCopyOfTransferFlags),
 			IsoNumDesc:             uint32(usb.IsoNumDesc),
+			Payload:                payload,
 		}
 	}
 	return nil

diff --git a/netcap.proto b/netcap.proto
@@ -547,6 +547,7 @@ message UDP {
     int32  Checksum       = 5;
     double PayloadEntropy = 6;
     int32  PayloadSize    = 7;
+    bytes  Payload        = 8;
 }
 
 message TCP {
@@ -573,6 +574,7 @@ message TCP {
     repeated TCPOption Options = 20;
     double   PayloadEntropy    = 21;
     int32    PayloadSize       = 22;
+    bytes    Payload           = 23;
 }
 
 message TCPOption {
@@ -948,6 +950,7 @@ message USB {
     uint32      UrbStartFrame             = 17;
     uint32      UrbCopyOfTransferFlags    = 18;
     uint32      IsoNumDesc                = 19;
+    bytes       Payload                   = 20;
 }
 
 message USBRequestBlockSetup {