From 94b834936ad223646a57399dd9f90ef081446a5f Mon Sep 17 00:00:00 2001
From: Chaz Lever <chazlever@users.noreply.github.com>
Date: Fri, 22 Nov 2024 16:04:06 -0500
Subject: [PATCH] Add --bpf-filter flag to support custom BPF filters

This commit removes the `--tcp` flag, removes the default BPF filter,
and adds support for supplying custom BPF filters through the
`--bpf-filter` flag.
---
 main.go         |  8 +++----
 parser/parse.go | 64 ++++++++++++++++++++++---------------------------
 2 files changed, 32 insertions(+), 40 deletions(-)

diff --git a/main.go b/main.go
index 91d2ae8..f2b79bf 100644
--- a/main.go
+++ b/main.go
@@ -22,7 +22,7 @@ func getOutputFormats() []string {
 }
 
 func loadGlobalOptions(c *cli.Context) error {
-	parser.DoParseTcp = c.GlobalBool("tcp")
+	parser.BpfFilter = c.GlobalString("bpf-filter")
 	parser.DoParseQuestions = c.GlobalBool("questions")
 	parser.DoParseQuestionsEcs = c.GlobalBool("questions-ecs")
 	parser.Source = c.GlobalString("source")
@@ -133,9 +133,9 @@ func main() {
 	}
 
 	app.Flags = []cli.Flag{
-		cli.BoolFlag{
-			Name:  "tcp",
-			Usage: "attempt to parse TCP packets",
+		cli.StringFlag{
+			Name:  "bpf-filter",
+			Usage: "specify a BPF filter to use for filtering packets",
 		},
 		cli.BoolFlag{
 			Name:  "questions",
diff --git a/parser/parse.go b/parser/parse.go
index ac2fb61..5685ca3 100644
--- a/parser/parse.go
+++ b/parser/parse.go
@@ -16,7 +16,7 @@ import (
 )
 
 var (
-	DoParseTcp          = true
+	BpfFilter           = ""
 	DoParseQuestions    = false
 	DoParseQuestionsEcs = true
 	Source              = ""
@@ -42,15 +42,12 @@ func ParseFile(fname string) {
 	defer handle.Close()
 
 	// Setup BPF filter on handle
-	bpfFilter := "udp port 53 or (vlan and udp port 53)"
-	if DoParseTcp {
-		bpfFilter = "port 53 or (vlan and port 53)"
-	}
-	err = handle.SetBPFFilter(bpfFilter)
-	if err != nil {
-		log.Warnf("Could not set BPF filter: %v\n", err)
+	if BpfFilter != "" {
+		err = handle.SetBPFFilter(BpfFilter)
+		if err != nil {
+			log.Warnf("Could not set BPF filter: %v\n", err)
+		}
 	}
-
 	ParseDns(handle)
 }
 
@@ -62,15 +59,12 @@ func ParseDevice(device string, snapshotLen int32, promiscuous bool, timeout tim
 	defer handle.Close()
 
 	// Setup BPF filter on handle
-	bpfFilter := "udp port 53 or (vlan and udp port 53)"
-	if DoParseTcp {
-		bpfFilter = "port 53 or (vlan and port 53)"
-	}
-	err = handle.SetBPFFilter(bpfFilter)
-	if err != nil {
-		log.Warnf("Could not set BPF filter: %v\n", err)
+	if BpfFilter != "" {
+		err = handle.SetBPFFilter(BpfFilter)
+		if err != nil {
+			log.Warnf("Could not set BPF filter: %v\n", err)
+		}
 	}
-
 	ParseDns(handle)
 }
 
@@ -80,7 +74,7 @@ func ParseDns(handle *pcap.Handle) {
 		stats  Statistics
 		ip4    *layers.IPv4
 		ip6    *layers.IPv6
-		tcp    *layers.TCP
+		_      *layers.TCP
 		udp    *layers.UDP
 		msg    *dns.Msg
 	)
@@ -143,25 +137,23 @@ PACKETLOOP:
 		}
 		switch transportLayer.LayerType() {
 		case layers.LayerTypeTCP:
-			tcp = transportLayer.(*layers.TCP)
 			stats.PacketTcp += 1
-
-			if !DoParseTcp {
-				continue PACKETLOOP
-			}
-
-			msg = new(dns.Msg)
-			if err := msg.Unpack(tcp.Payload); err != nil {
-				log.Errorf("Could not decode DNS: %v\n", err)
-				stats.PacketErrors += 1
-				continue PACKETLOOP
-			}
-			stats.PacketDns += 1
-
-			schema.SourcePort = uint16(tcp.SrcPort)
-			schema.DestinationPort = uint16(tcp.DstPort)
-			schema.Udp = false
-			schema.Sha256 = fmt.Sprintf("%x", sha256.Sum256(tcp.Payload))
+			continue PACKETLOOP
+			// TODO: Implement TCP reassembly for DNS parsing
+			//tcp = transportLayer.(*layers.TCP)
+			//
+			//msg = new(dns.Msg)
+			//if err := msg.Unpack(tcp.Payload); err != nil {
+			//	log.Errorf("Could not decode DNS: %v\n", err)
+			//	stats.PacketErrors += 1
+			//	continue PACKETLOOP
+			//}
+			//stats.PacketDns += 1
+			//
+			//schema.SourcePort = uint16(tcp.SrcPort)
+			//schema.DestinationPort = uint16(tcp.DstPort)
+			//schema.Udp = false
+			//schema.Sha256 = fmt.Sprintf("%x", sha256.Sum256(tcp.Payload))
 		case layers.LayerTypeUDP:
 			udp = transportLayer.(*layers.UDP)
 			stats.PacketUdp += 1