misc: small fix or general refactoring i did not bother commenting

modules/net_probe_udp.go

@@ -3,6 +3,7 @@ package modules
 import (
 	"fmt"
 	"net"
+	"sync/atomic"
 
 	"github.com/evilsocket/bettercap-ng/log"
 )
@@ -19,9 +20,7 @@ func (p *Prober) sendProbeUDP(from net.IP, from_hw net.HardwareAddr, ip net.IP)
 		wrote, _ := con.Write([]byte{0x00})
 
 		if wrote > 0 {
-			p.Session.Queue.Lock()
-			p.Session.Queue.Sent += uint64(wrote)
-			p.Session.Queue.Unlock()
+			atomic.AddUint64(&p.Session.Queue.Stats.Sent, uint64(wrote))
 		}
 	}
 }

modules/net_recon_show.go

@@ -144,11 +144,11 @@ func (d *Discovery) Show(by string) error {
 
 	fmt.Printf("\n%s %s / %s %s / %d pkts / %d errs\n\n",
 		core.Red("↑"),
-		humanize.Bytes(atomic.LoadUint64(&d.Session.Queue.Sent)),
+		humanize.Bytes(atomic.LoadUint64(&d.Session.Queue.Stats.Sent)),
 		core.Green("↓"),
-		humanize.Bytes(atomic.LoadUint64(&d.Session.Queue.Received)),
-		atomic.LoadUint64(&d.Session.Queue.PktReceived),
-		atomic.LoadUint64(&d.Session.Queue.Errors))
+		humanize.Bytes(atomic.LoadUint64(&d.Session.Queue.Stats.Received)),
+		atomic.LoadUint64(&d.Session.Queue.Stats.PktReceived),
+		atomic.LoadUint64(&d.Session.Queue.Stats.Errors))
 
 	s := EventsStream{}
 	events := d.Session.Events.Sorted()

packets/queue.go

@@ -25,15 +25,19 @@ type Traffic struct {
 	Received uint64
 }
 
-type Queue struct {
-	sync.Mutex
-
+type Stats struct {
 	Sent        uint64
 	Received    uint64
 	PktReceived uint64
 	Errors      uint64
+}
+
+type Queue struct {
+	sync.Mutex
 
 	Activities chan Activity `json:"-"`
+
+	Stats   Stats
 	Protos  map[string]uint64
 	Traffic map[string]*Traffic
 
@@ -43,32 +47,24 @@ type Queue struct {
 	active bool
 }
 
-func NewQueue(iface *bnet.Endpoint) (*Queue, error) {
-	var err error
-
-	q := &Queue{
-		iface:       iface,
-		handle:      nil,
-		active:      true,
-		source:      nil,
-		Sent:        0,
-		Received:    0,
-		PktReceived: 0,
-		Errors:      0,
+func NewQueue(iface *bnet.Endpoint) (q *Queue, err error) {
+	q = &Queue{
 		Protos:     make(map[string]uint64),
 		Traffic:    make(map[string]*Traffic),
 		Activities: make(chan Activity),
+
+		iface:  iface,
+		active: true,
 	}
 
-	q.handle, err = pcap.OpenLive(iface.Name(), 1024, true, pcap.BlockForever)
-	if err != nil {
-		return nil, err
+	if q.handle, err = pcap.OpenLive(iface.Name(), 1024, true, pcap.BlockForever); err != nil {
+		return
 	}
 
 	q.source = gopacket.NewPacketSource(q.handle, q.handle.LinkType())
 	go q.worker()
 
-	return q, nil
+	return
 }
 
 func (q *Queue) trackProtocols(pkt gopacket.Packet) {
@@ -90,9 +86,7 @@ func (q *Queue) trackProtocols(pkt gopacket.Packet) {
 	}
 }
 
-func (q *Queue) trackActivity(eth *layers.Ethernet, ip4 *layers.IPv4, address net.IP, pktSize uint64) {
-	// detrmine direction
-	isSent := bytes.Compare(address, ip4.SrcIP) == 0
+func (q *Queue) trackActivity(eth *layers.Ethernet, ip4 *layers.IPv4, address net.IP, pktSize uint64, isSent bool) {
 	// push to activity channel
 	q.Activities <- Activity{
 		IP:     address,
@@ -130,8 +124,8 @@ func (q *Queue) worker() {
 
 		pktSize := uint64(len(pkt.Data()))
 
-		atomic.AddUint64(&q.PktReceived, 1)
-		atomic.AddUint64(&q.Received, pktSize)
+		atomic.AddUint64(&q.Stats.PktReceived, 1)
+		atomic.AddUint64(&q.Stats.Received, pktSize)
 
 		// decode eth and ipv4 layers
 		leth := pkt.Layer(layers.LayerTypeEthernet)
@@ -142,11 +136,11 @@ func (q *Queue) worker() {
 
 			// coming from our network
 			if bytes.Compare(q.iface.IP, ip4.SrcIP) != 0 && q.iface.Net.Contains(ip4.SrcIP) {
-				q.trackActivity(eth, ip4, ip4.SrcIP, pktSize)
+				q.trackActivity(eth, ip4, ip4.SrcIP, pktSize, true)
 			}
 			// coming to our network
 			if bytes.Compare(q.iface.IP, ip4.DstIP) != 0 && q.iface.Net.Contains(ip4.DstIP) {
-				q.trackActivity(eth, ip4, ip4.DstIP, pktSize)
+				q.trackActivity(eth, ip4, ip4.DstIP, pktSize, false)
 			}
 		}
 	}
@@ -161,10 +155,10 @@ func (q *Queue) Send(raw []byte) error {
 	}
 
 	if err := q.handle.WritePacketData(raw); err != nil {
-		q.Errors += 1
+		atomic.AddUint64(&q.Stats.Errors, 1)
 		return err
 	} else {
-		q.Sent += uint64(len(raw))
+		atomic.AddUint64(&q.Stats.Sent, uint64(len(raw)))
 	}
 
 	return nil

session/prompt.go

@@ -37,22 +37,22 @@ var PromptCallbacks = map[string]func(s *Session) string{
 		return s.Interface.CIDR()
 	},
 	"{net.sent}": func(s *Session) string {
-		return fmt.Sprintf("%d", s.Queue.Sent)
+		return fmt.Sprintf("%d", s.Queue.Stats.Sent)
 	},
 	"{net.sent.human}": func(s *Session) string {
-		return humanize.Bytes(s.Queue.Sent)
+		return humanize.Bytes(s.Queue.Stats.Sent)
 	},
 	"{net.received}": func(s *Session) string {
-		return fmt.Sprintf("%d", s.Queue.Received)
+		return fmt.Sprintf("%d", s.Queue.Stats.Received)
 	},
 	"{net.received.human}": func(s *Session) string {
-		return humanize.Bytes(s.Queue.Received)
+		return humanize.Bytes(s.Queue.Stats.Received)
 	},
 	"{net.packets}": func(s *Session) string {
-		return fmt.Sprintf("%d", s.Queue.PktReceived)
+		return fmt.Sprintf("%d", s.Queue.Stats.PktReceived)
 	},
 	"{net.errors}": func(s *Session) string {
-		return fmt.Sprintf("%d", s.Queue.Errors)
+		return fmt.Sprintf("%d", s.Queue.Stats.Errors)
 	},
 }