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

This commit is contained in:
evilsocket 2018-02-11 19:43:36 +01:00
parent ae81bb8606
commit ea477ef2b4
4 changed files with 38 additions and 45 deletions

View file

@ -3,6 +3,7 @@ package modules
import ( import (
"fmt" "fmt"
"net" "net"
"sync/atomic"
"github.com/evilsocket/bettercap-ng/log" "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}) wrote, _ := con.Write([]byte{0x00})
if wrote > 0 { if wrote > 0 {
p.Session.Queue.Lock() atomic.AddUint64(&p.Session.Queue.Stats.Sent, uint64(wrote))
p.Session.Queue.Sent += uint64(wrote)
p.Session.Queue.Unlock()
} }
} }
} }

View file

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

View file

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

View file

@ -37,22 +37,22 @@ var PromptCallbacks = map[string]func(s *Session) string{
return s.Interface.CIDR() return s.Interface.CIDR()
}, },
"{net.sent}": func(s *Session) string { "{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 { "{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 { "{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 { "{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 { "{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 { "{net.errors}": func(s *Session) string {
return fmt.Sprintf("%d", s.Queue.Errors) return fmt.Sprintf("%d", s.Queue.Stats.Errors)
}, },
} }