diff --git a/modules/net_probe_udp.go b/modules/net_probe_udp.go index 4d14df7b..904ec6c6 100644 --- a/modules/net_probe_udp.go +++ b/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)) } } } diff --git a/modules/net_recon_show.go b/modules/net_recon_show.go index c9de6590..5206cbd5 100644 --- a/modules/net_recon_show.go +++ b/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() diff --git a/packets/queue.go b/packets/queue.go index 25141952..c4b755d1 100644 --- a/packets/queue.go +++ b/packets/queue.go @@ -25,17 +25,21 @@ 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:"-"` - Protos map[string]uint64 - Traffic map[string]*Traffic + + Stats Stats + Protos map[string]uint64 + Traffic map[string]*Traffic iface *bnet.Endpoint handle *pcap.Handle @@ -43,32 +47,24 @@ type Queue struct { active bool } -func NewQueue(iface *bnet.Endpoint) (*Queue, error) { - var err error +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), - q := &Queue{ - iface: iface, - handle: nil, - active: true, - source: nil, - Sent: 0, - Received: 0, - PktReceived: 0, - Errors: 0, - 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 diff --git a/session/prompt.go b/session/prompt.go index a5e33bd0..708b2a59 100644 --- a/session/prompt.go +++ b/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) }, }