Small refactoring

Author: shadowy-pycoder

mshark.go

@@ -46,13 +46,19 @@ type Writer struct {
 	packetNum 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"))
@@ -115,6 +121,19 @@ func (mw *Writer) WritePacket(timestamp time.Time, data []byte) error {
 	return nil
 }
 
+// WriteHeader writes a header to the writer.
+//
+// The header contains metadata about the capture, such as the interface name,
+// snapshot length, promiscuous mode, timeout, number of packets, and BPF filter.
+//
+// The header is written in the following format:
+//
+//   - Interface: eth0
+//   - Snapshot Length: 65535
+//   - Promiscuous Mode: true
+//   - Timeout: 5s
+//   - Number of Packets: 0
+//   - BPF Filter: "ip proto tcp"
 func (mw *Writer) WriteHeader(c *Config, in *net.Interface, snaplen int) error {
 	_, err := fmt.Fprintf(mw.w, `- Interface: %s
 - Snapshot Length: %d
@@ -134,6 +153,31 @@ func (mw *Writer) WriteHeader(c *Config, in *net.Interface, snaplen int) error {
 	return err
 }
 
+// InterfaceByName returns the interface specified by name.
+func InterfaceByName(name string) (*net.Interface, error) {
+	var (
+		in  *net.Interface
+		err error
+	)
+	if name == "any" {
+		in = &net.Interface{Index: 0, Name: "any"}
+	} else {
+		in, err = net.InterfaceByName(name)
+		if err != nil {
+			return nil, fmt.Errorf("unknown interface %s: %v", name, err)
+		}
+		ok := true &&
+			// Look for an Ethernet interface.
+			len(in.HardwareAddr) == 6 &&
+			// Look for up, multicast, broadcast.
+			in.Flags&(net.FlagUp|net.FlagMulticast|net.FlagBroadcast) != 0
+		if !ok {
+			return nil, fmt.Errorf("interface %s is not up", name)
+		}
+	}
+	return in, nil
+}
+
 func OpenLive(conf *Config) error {
 
 	packetcfg := packet.Config{}
@@ -153,26 +197,10 @@ func OpenLive(conf *Config) error {
 		packetcfg.Filter = raw
 	}
 
-	var (
-		in  *net.Interface
-		err error
-	)
-	iface := conf.Iface
-	if iface == "any" {
-		in = &net.Interface{Index: 0, Name: "any"}
-	} else {
-		in, err = net.InterfaceByName(iface)
-		if err != nil {
-			return fmt.Errorf("unknown interface %s: %v", iface, err)
-		}
-		ok := true &&
-			// Look for an Ethernet interface.
-			len(in.HardwareAddr) == 6 &&
-			// Look for up, multicast, broadcast.
-			in.Flags&(net.FlagUp|net.FlagMulticast|net.FlagBroadcast) != 0
-		if !ok {
-			return fmt.Errorf("interface %s is not up", iface)
-		}
+	// getting interface
+	in, err := InterfaceByName(conf.Iface)
+	if err != nil {
+		return err
 	}
 
 	// opening connection
@@ -186,14 +214,14 @@ func OpenLive(conf *Config) error {
 
 	// setting promisc mode
 	if in.Name != "any" {
-		if err = c.SetPromiscuous(conf.Promisc); err != nil {
+		if err := c.SetPromiscuous(conf.Promisc); err != nil {
 			return fmt.Errorf("unable to set promiscuous mode: %v", err)
 		}
 	}
 
 	// timeout
 	if conf.Timeout > 0 {
-		if err = c.SetDeadline(time.Now().Add(conf.Timeout)); err != nil {
+		if err := c.SetDeadline(time.Now().Add(conf.Timeout)); err != nil {
 			return fmt.Errorf("unable to set timeout: %v", err)
 		}
 	}
@@ -206,31 +234,9 @@ func OpenLive(conf *Config) error {
 
 	// file to write packets
 	f := NewWriter(conf.File)
-	if err = f.WriteHeader(conf, in, snaplen); err != nil {
+	if err := f.WriteHeader(conf, in, snaplen); err != nil {
 		return err
 	}
-	// number of packets
-	count := conf.PacketCount
-	if count < 0 {
-		count = 0
-	}
-	infinity := count == 0
-
-	defer func() {
-		stats, err := c.Stats()
-		if err != nil {
-			fmt.Printf("failed to fetch stats: %v", err)
-		}
-		// Received some data, assume some Stats were populated.
-		if stats.Packets == 0 {
-			fmt.Println("stats indicated 0 received packets")
-		}
-
-		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)
-		// close Conn
-		c.Close()
-	}()
 
 	// pcap file to write packets
 	var pcap *mpcap.Writer
@@ -242,7 +248,7 @@ func OpenLive(conf *Config) error {
 		}
 		defer fpcap.Close()
 		pcap = mpcap.NewWriter(fpcap)
-		if err = pcap.WriteHeader(snaplen); err != nil {
+		if err := pcap.WriteHeader(snaplen); err != nil {
 			return err
 		}
 	}
@@ -257,11 +263,30 @@ func OpenLive(conf *Config) error {
 		}
 		defer fpcapng.Close()
 		pcapng = mpcapng.NewWriter(fpcapng)
-		if err = pcapng.WriteHeader("mshark", in, conf.Expr, conf.Snaplen); err != nil {
+		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)
+		}
+		// close Conn
+		c.Close()
+	}()
+
+	// number of packets
+	count := conf.PacketCount
+	if count < 0 {
+		count = 0
+	}
+	infinity := count == 0
+
 	for i := 0; infinity || i < count; i++ {
 		n, _, err := c.ReadFrom(b)
 		if err != nil {
@@ -275,19 +300,19 @@ func OpenLive(conf *Config) error {
 		data := b[:n]
 
 		if conf.Pcap {
-			if err = pcap.WritePacket(timestamp, data); err != nil {
+			if err := pcap.WritePacket(timestamp, data); err != nil {
 				return err
 			}
 		}
 
 		if conf.PcapNG {
-			if err = pcapng.WritePacket(timestamp, data); err != nil {
+			if err := pcapng.WritePacket(timestamp, data); err != nil {
 				return err
 			}
 		}
 
-		if err = f.WritePacket(timestamp, data); err != nil {
-			fmt.Fprintln(os.Stderr, err)
+		if err := f.WritePacket(timestamp, data); err != nil {
+			return err
 		}
 	}
 	return nil