Made OpenLive accept custom packet writers

Author: shadowy-pycoder

.gitignore

@@ -28,4 +28,5 @@ tls/
 
 *.prof
 
-*.pcap*
+*.pcap*
+*.txt

cmd/mshark/cli.go

@@ -8,8 +8,11 @@ import (
 	"path/filepath"
 	"strings"
 	"text/tabwriter"
+	"time"
 
 	ms "github.com/shadowy-pycoder/mshark"
+	"github.com/shadowy-pycoder/mshark/mpcap"
+	"github.com/shadowy-pycoder/mshark/mpcapng"
 )
 
 const usagePrefix = `
@@ -47,28 +50,29 @@ func displayInterfaces() error {
 	return w.Flush()
 }
 
+func createFile(ext string) (*os.File, error) {
+	path := fmt.Sprintf("./mshark_%s.%s", time.Now().UTC().Format("20060102_150405"), ext)
+	f, err := os.OpenFile(filepath.FromSlash(path), os.O_CREATE|os.O_WRONLY, 0644)
+	if err != nil {
+		return nil, fmt.Errorf("failed to open file: %v", err)
+	}
+	return f, nil
+}
+
 func root(args []string) error {
 	conf := ms.Config{}
 
 	flags := flag.NewFlagSet("mshark", flag.ExitOnError)
-	flags.StringVar(&conf.Iface, "i", "any", "The name of the network interface. Example: eth0")
-	flags.IntVar(&conf.Snaplen, "s", 0, "The maximum length of each packet snapshot. Defaults to 65535.")
+	iface := flags.String("i", "any", "The name of the network interface. Example: eth0")
+	snaplen := flags.Int("s", 0, "The maximum length of each packet snapshot. Defaults to 65535.")
 	flags.BoolFunc("p", `Promiscuous mode. This setting is ignored for "any" interface. Defaults to false.`, func(flagValue string) error {
 		conf.Promisc = true
 		return nil
 	})
 	flags.DurationVar(&conf.Timeout, "t", 0, "The maximum duration of the packet capture process. Example: 5s")
 	flags.IntVar(&conf.PacketCount, "c", 0, "The maximum number of packets to capture.")
 	flags.StringVar(&conf.Expr, "e", "", `BPF filter expression. Example: "ip proto tcp"`)
-	path := flags.String("f", "", "File path to write captured packet data to. Defaults to stdout.")
-	flags.BoolFunc("pcap", "Create a PCAP file in the current working directory.", func(flagValue string) error {
-		conf.Pcap = true
-		return nil
-	})
-	flags.BoolFunc("pcapng", "Create a PCAPNG file in the current working directory.", func(flagValue string) error {
-		conf.PcapNG = true
-		return nil
-	})
+	ext := flags.String("f", "", "File extension to write captured packet data to. Supported values: txt, pcap, pcapng")
 	flags.BoolFunc("D", "Display list of interfaces and exit.", func(flagValue string) error {
 		if err := displayInterfaces(); err != nil {
 			fmt.Fprintf(os.Stderr, "mshark: %v\n", err)
@@ -87,16 +91,67 @@ func root(args []string) error {
 		return err
 	}
 
-	if *path != "" {
-		f, err := os.OpenFile(filepath.FromSlash(*path), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
-		if err != nil {
-			return fmt.Errorf("failed to open file: %v", err)
+	// getting network interface from the provided name
+	in, err := ms.InterfaceByName(*iface)
+	if err != nil {
+		return err
+	}
+	conf.Device = in
+
+	// checking snaplen
+	if *snaplen <= 0 || *snaplen > 65535 {
+		*snaplen = 65535
+	}
+	conf.Snaplen = *snaplen
+
+	// creating writers depending on a file extension
+	var pw ms.PacketWriter
+	if *ext != "" {
+		switch *ext {
+		case "txt":
+			f, err := createFile(*ext)
+			if err != nil {
+				return err
+			}
+			defer f.Close()
+			pw = ms.NewWriter(f)
+		case "pcap":
+			f, err := createFile(*ext)
+			if err != nil {
+				return err
+			}
+			defer f.Close()
+			pw = mpcap.NewWriter(f)
+		case "pcapng":
+			f, err := createFile(*ext)
+			if err != nil {
+				return err
+			}
+			defer f.Close()
+			pw = mpcapng.NewWriter(f)
+		default:
+			return fmt.Errorf("unsupported file format: %s", *ext)
 		}
-		defer f.Close()
-		conf.File = f
+	} else {
+		pw = ms.NewWriter(os.Stdout)
 	}
 
-	if err := ms.OpenLive(&conf); err != nil {
+	// writing headers
+	switch t := pw.(type) {
+	case *ms.Writer:
+		if err := t.WriteHeader(&conf); err != nil {
+			return err
+		}
+	case *mpcap.Writer:
+		if err := t.WriteHeader(conf.Snaplen); err != nil {
+			return err
+		}
+	case *mpcapng.Writer:
+		if err := t.WriteHeader("mshark", conf.Device, conf.Expr, conf.Snaplen); err != nil {
+			return err
+		}
+	}
+	if err := ms.OpenLive(&conf, pw); err != nil {
 		return err
 	}
 	return nil

mshark.go

@@ -6,17 +6,18 @@ import (
 	"io"
 	"net"
 	"os"
-	"path/filepath"
 	"time"
 
 	"github.com/mdlayher/packet"
 	"github.com/packetcap/go-pcap/filter"
-	"github.com/shadowy-pycoder/mshark/mpcap"
-	"github.com/shadowy-pycoder/mshark/mpcapng"
 	"golang.org/x/net/bpf"
 	"golang.org/x/sys/unix"
 )
 
+type PacketWriter interface {
+	WritePacket(timestamp time.Time, data []byte) error
+}
+
 type layers struct {
 	ether EthernetFrame
 	ip    IPv4Packet
@@ -29,39 +30,31 @@ type layers struct {
 }
 
 type Config struct {
-	Iface       string        // The name of the network interface ("any" means listen on all interfaces).
-	Snaplen     int           // The maximum length of each packet snapshot.
-	Promisc     bool          // Promiscuous mode. This setting is ignored for "any" interface.
-	Timeout     time.Duration // The maximum duration of the packet capture process.
-	PacketCount int           // The maximum number of packets to capture.
-	Expr        string        // BPF filter expression.
-	File        io.Writer     // File to write captured packet data to.
-	Pcap        bool          // Create a PCAP file in the current working directory.
-	PcapNG      bool          // Create a PCAPNG file in the current working directory.
+	Device      *net.Interface // The name of the network interface ("any" means listen on all interfaces).
+	Snaplen     int            // The maximum length of each packet snapshot.
+	Promisc     bool           // Promiscuous mode. This setting is ignored for "any" interface.
+	Timeout     time.Duration  // The maximum duration of the packet capture process.
+	PacketCount int            // The maximum number of packets to capture.
+	Expr        string         // BPF filter expression.
 }
 
 type Writer struct {
-	w         io.Writer
-	layers    layers
-	packetNum uint64
+	w       io.Writer
+	layers  layers
+	packets uint64
 }
 
 // NewWriter creates a new mshark Writer.
-//
-// If the io.Writer is nil, the Writer will write to stdout.
 func NewWriter(w io.Writer) *Writer {
-	if w == nil {
-		w = os.Stdout
-	}
 	return &Writer{w: w}
 }
 
 // WritePacket writes a packet to the writer, along with its timestamp.
 //
 // Timestamps are to be generated by the calling code.
 func (mw *Writer) WritePacket(timestamp time.Time, data []byte) error {
-	mw.packetNum++
-	fmt.Fprintf(mw.w, "- Packet: %d Timestamp: %s\n", mw.packetNum, timestamp.Format("2006-01-02T15:04:05-0700"))
+	mw.packets++
+	fmt.Fprintf(mw.w, "- Packet: %d Timestamp: %s\n", mw.packets, timestamp.Format("2006-01-02T15:04:05-0700"))
 	fmt.Fprintln(mw.w, "==================================================================")
 	if err := mw.layers.ether.Parse(data); err != nil {
 		return err
@@ -134,7 +127,7 @@ func (mw *Writer) WritePacket(timestamp time.Time, data []byte) error {
 //   - Timeout: 5s
 //   - Number of Packets: 0
 //   - BPF Filter: "ip proto tcp"
-func (mw *Writer) WriteHeader(c *Config, in *net.Interface, snaplen int) error {
+func (mw *Writer) WriteHeader(c *Config) error {
 	_, err := fmt.Fprintf(mw.w, `- Interface: %s
 - Snapshot Length: %d
 - Promiscuous Mode: %v
@@ -143,9 +136,9 @@ func (mw *Writer) WriteHeader(c *Config, in *net.Interface, snaplen int) error {
 - BPF Filter: %q
 
 `,
-		in.Name,
-		snaplen,
-		in.Name != "any" && c.Promisc,
+		c.Device.Name,
+		c.Snaplen,
+		c.Device.Name != "any" && c.Promisc,
 		c.Timeout,
 		c.PacketCount,
 		c.Expr,
@@ -178,7 +171,7 @@ func InterfaceByName(name string) (*net.Interface, error) {
 	return in, nil
 }
 
-func OpenLive(conf *Config) error {
+func OpenLive(conf *Config, pw PacketWriter) error {
 
 	packetcfg := packet.Config{}
 
@@ -197,14 +190,8 @@ func OpenLive(conf *Config) error {
 		packetcfg.Filter = raw
 	}
 
-	// getting interface
-	in, err := InterfaceByName(conf.Iface)
-	if err != nil {
-		return err
-	}
-
 	// opening connection
-	c, err := packet.Listen(in, packet.Raw, unix.ETH_P_ALL, &packetcfg)
+	c, err := packet.Listen(conf.Device, packet.Raw, unix.ETH_P_ALL, &packetcfg)
 	if err != nil {
 		if errors.Is(err, os.ErrPermission) {
 			return fmt.Errorf("permission denied (try setting CAP_NET_RAW capability): %v", err)
@@ -213,7 +200,7 @@ func OpenLive(conf *Config) error {
 	}
 
 	// setting promisc mode
-	if in.Name != "any" {
+	if conf.Device.Name != "any" {
 		if err := c.SetPromiscuous(conf.Promisc); err != nil {
 			return fmt.Errorf("unable to set promiscuous mode: %v", err)
 		}
@@ -225,56 +212,17 @@ func OpenLive(conf *Config) error {
 			return fmt.Errorf("unable to set timeout: %v", err)
 		}
 	}
-	// snaplen
-	snaplen := conf.Snaplen
-	if snaplen <= 0 || snaplen > 65535 {
-		snaplen = 65535
-	}
-	b := make([]byte, snaplen)
-
-	// file to write packets
-	f := NewWriter(conf.File)
-	if err := f.WriteHeader(conf, in, snaplen); err != nil {
-		return err
-	}
-
-	// pcap file to write packets
-	var pcap *mpcap.Writer
-	if conf.Pcap {
-		path := fmt.Sprintf("./mshark_%s.pcap", time.Now().UTC().Format("20060102_150405"))
-		fpcap, err := os.OpenFile(filepath.FromSlash(path), os.O_CREATE|os.O_WRONLY, 0644)
-		if err != nil {
-			return fmt.Errorf("failed to open file: %v", err)
-		}
-		defer fpcap.Close()
-		pcap = mpcap.NewWriter(fpcap)
-		if err := pcap.WriteHeader(snaplen); err != nil {
-			return err
-		}
-	}
-
-	// pcapng file to write packets
-	var pcapng *mpcapng.Writer
-	if conf.PcapNG {
-		path := fmt.Sprintf("./mshark_%s.pcapng", time.Now().UTC().Format("20060102_150405"))
-		fpcapng, err := os.OpenFile(filepath.FromSlash(path), os.O_CREATE|os.O_WRONLY, 0644)
-		if err != nil {
-			return fmt.Errorf("failed to open file: %v", err)
-		}
-		defer fpcapng.Close()
-		pcapng = mpcapng.NewWriter(fpcapng)
-		if err := pcapng.WriteHeader("mshark", in, conf.Expr, snaplen); err != nil {
-			return err
-		}
-	}
 
 	defer func() {
 		stats, err := c.Stats()
 		if err != nil {
 			fmt.Fprintf(os.Stderr, "failed to fetch stats: %v", err)
 		} else {
-			fmt.Fprintf(f.w, "- Packets: %d, Drops: %d, Freeze Queue Count: %d\n- Packets Captured: %d\n",
-				stats.Packets, stats.Drops, stats.FreezeQueueCount, f.packetNum)
+			fmt.Printf("- Packets: %d, Drops: %d, Freeze Queue Count: %d\n",
+				stats.Packets, stats.Drops, stats.FreezeQueueCount)
+			if w, ok := pw.(*Writer); ok {
+				fmt.Printf("- Packets Captured: %d\n", w.packets)
+			}
 		}
 		// close Conn
 		c.Close()
@@ -287,6 +235,8 @@ func OpenLive(conf *Config) error {
 	}
 	infinity := count == 0
 
+	b := make([]byte, conf.Snaplen)
+
 	for i := 0; infinity || i < count; i++ {
 		n, _, err := c.ReadFrom(b)
 		if err != nil {
@@ -296,22 +246,7 @@ func OpenLive(conf *Config) error {
 			return fmt.Errorf("failed to read Ethernet frame: %v", err)
 		}
 
-		timestamp := time.Now().UTC()
-		data := b[:n]
-
-		if conf.Pcap {
-			if err := pcap.WritePacket(timestamp, data); err != nil {
-				return err
-			}
-		}
-
-		if conf.PcapNG {
-			if err := pcapng.WritePacket(timestamp, data); err != nil {
-				return err
-			}
-		}
-
-		if err := f.WritePacket(timestamp, data); err != nil {
+		if err := pw.WritePacket(time.Now().UTC(), b[:n]); err != nil {
 			return err
 		}
 	}

mshark_test.go

@@ -7,14 +7,18 @@ import (
 
 func BenchmarkOpenLive(b *testing.B) {
 	b.ResetTimer()
+	in, err := InterfaceByName("any")
+	if err != nil {
+		b.Fatal(err)
+	}
 	conf := Config{
-		Iface:       "any",
+		Device:      in,
 		Snaplen:     1600,
 		Promisc:     true,
 		PacketCount: b.N,
-		File:        io.Discard,
 	}
-	if err := OpenLive(&conf); err != nil {
+	pw := NewWriter(io.Discard)
+	if err := OpenLive(&conf, pw); err != nil {
 		b.Fatal(err)
 	}
 }