Merge remote-tracking branch 'upstream/master'

Makefile

@@ -6,6 +6,9 @@ all: deps build
 deps: godep golint gofmt gomegacheck
 	@dep ensure
 
+build_with_race_detector: resources
+	@go build -race -o $(TARGET) .
+
 build: resources
 	@go build -o $(TARGET) .
 

modules/ble/ble_show_services.go

@@ -367,6 +367,10 @@ func (mod *BLERecon) showServices(p gatt.Peripheral, services []*gatt.Service) {
 					data = parseRawData(raw)
 				}
 
+				if ch.Name() == "Device Name" && data != "" && mod.currDevice.DeviceName == "" {
+					mod.currDevice.DeviceName = data
+				}
+
 				if multi == nil {
 					char.Data = data
 					rows = append(rows, []string{

modules/net_recon/net_show.go

@@ -61,8 +61,11 @@ func (mod *Discovery) getRow(e *network.Endpoint, withMeta bool) [][]string {
 
 	var traffic *packets.Traffic
 	var found bool
-	if traffic, found = mod.Session.Queue.Traffic[e.IpAddress]; !found {
+	var v interface{}
+	if v, found = mod.Session.Queue.Traffic.Load(e.IpAddress); !found {
 		traffic = &packets.Traffic{}
+	} else {
+		traffic = v.(*packets.Traffic)
 	}
 
 	seen := e.LastSeen.Format("15:04:05")
@@ -203,9 +206,6 @@ func (mod *Discovery) colNames(hasMeta bool) []string {
 }
 
 func (mod *Discovery) showStatusBar() {
-	mod.Session.Queue.Stats.RLock()
-	defer mod.Session.Queue.Stats.RUnlock()
-
 	parts := []string{
 		fmt.Sprintf("%s %s", tui.Red("↑"), humanize.Bytes(mod.Session.Queue.Stats.Sent)),
 		fmt.Sprintf("%s %s", tui.Green("↓"), humanize.Bytes(mod.Session.Queue.Stats.Received)),

modules/net_recon/net_show_sort.go

@@ -41,24 +41,19 @@ func (a BySeenSorter) Less(i, j int) bool { return a[i].LastSeen.Before(a[j].Las
 
 type BySentSorter []*network.Endpoint
 
+func trafficOf(ip string) *packets.Traffic {
+	if v, found := session.I.Queue.Traffic.Load(ip); !found {
+		return &packets.Traffic{}
+	} else {
+		return v.(*packets.Traffic)
+	}
+}
+
 func (a BySentSorter) Len() int      { return len(a) }
 func (a BySentSorter) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
 func (a BySentSorter) Less(i, j int) bool {
-	session.I.Queue.Lock()
+	aTraffic := trafficOf(a[i].IpAddress)
-	defer session.I.Queue.Unlock()
+	bTraffic := trafficOf(a[j].IpAddress)
-
-	var found bool = false
-	var aTraffic *packets.Traffic = nil
-	var bTraffic *packets.Traffic = nil
-
-	if aTraffic, found = session.I.Queue.Traffic[a[i].IpAddress]; !found {
-		aTraffic = &packets.Traffic{}
-	}
-
-	if bTraffic, found = session.I.Queue.Traffic[a[j].IpAddress]; !found {
-		bTraffic = &packets.Traffic{}
-	}
-
 	return bTraffic.Sent > aTraffic.Sent
 }
 
@@ -67,20 +62,7 @@ type ByRcvdSorter []*network.Endpoint
 func (a ByRcvdSorter) Len() int      { return len(a) }
 func (a ByRcvdSorter) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
 func (a ByRcvdSorter) Less(i, j int) bool {
-	session.I.Queue.Lock()
+	aTraffic := trafficOf(a[i].IpAddress)
-	defer session.I.Queue.Unlock()
+	bTraffic := trafficOf(a[j].IpAddress)
-
-	var found bool = false
-	var aTraffic *packets.Traffic = nil
-	var bTraffic *packets.Traffic = nil
-
-	if aTraffic, found = session.I.Queue.Traffic[a[i].IpAddress]; !found {
-		aTraffic = &packets.Traffic{}
-	}
-
-	if bTraffic, found = session.I.Queue.Traffic[a[j].IpAddress]; !found {
-		bTraffic = &packets.Traffic{}
-	}
-
 	return bTraffic.Received > aTraffic.Received
 }

modules/wifi/wifi_recon.go

@@ -124,7 +124,7 @@ func (mod *WiFiModule) discoverClients(radiotap *layers.RadioTap, dot11 *layers.
 			freq := int(radiotap.ChannelFrequency)
 			rssi := radiotap.DBMAntennaSignal
 
-			if station, isNew := ap.AddClientIfNew(bssid, freq, rssi); isNew {
+			if station, isNew := ap.AddClientIfNew(bssid, freq, rssi, mod.Session.Lan.Aliases()); isNew {
 				mod.Session.Events.Add("wifi.client.new", ClientEvent{
 					AP:     ap,
 					Client: station,

modules/wifi/wifi_recon_handshakes.go

@@ -36,7 +36,7 @@ func (mod *WiFiModule) discoverHandshakes(radiotap *layers.RadioTap, dot11 *laye
 		staIsUs := bytes.Equal(staMac, mod.iface.HW)
 		station, found := ap.Get(staMac.String())
 		if !found {
-			station, _ = ap.AddClientIfNew(staMac.String(), ap.Frequency, ap.RSSI)
+			station, _ = ap.AddClientIfNew(staMac.String(), ap.Frequency, ap.RSSI, mod.Session.Lan.Aliases())
 		}
 
 		rawPMKID := []byte(nil)

modules/wifi/wifi_show.go

@@ -296,9 +296,6 @@ func (mod *WiFiModule) colNames(nrows int) []string {
 }
 
 func (mod *WiFiModule) showStatusBar() {
-	mod.Session.Queue.Stats.RLock()
-	defer mod.Session.Queue.Stats.RUnlock()
-
 	parts := []string{
 		fmt.Sprintf("%s (ch. %d)", mod.iface.Name(), network.GetInterfaceChannel(mod.iface.Name())),
 		fmt.Sprintf("%s %s", tui.Red("↑"), humanize.Bytes(mod.Session.Queue.Stats.Sent)),

network/ble_device.go

@@ -28,6 +28,7 @@ type BLEService struct {
 
 type BLEDevice struct {
 	LastSeen      time.Time
+	DeviceName    string
 	Vendor        string
 	RSSI          int
 	Device        gatt.Peripheral
@@ -62,10 +63,16 @@ func NewBLEDevice(p gatt.Peripheral, a *gatt.Advertisement, rssi int) *BLEDevice
 }
 
 func (d *BLEDevice) Name() string {
+	// get the name if it's being set during services enumeration via 'Device Name'
+	name := d.DeviceName
+	if name == "" {
+		// get the name from the device
 		name := d.Device.Name()
 		if name == "" && d.Advertisement != nil {
+			// get the name from the advertisement data
 			name = d.Advertisement.LocalName
 		}
+	}
 	return name
 }
 

network/lan.go

@@ -68,8 +68,8 @@ func (lan *LAN) SetAliasFor(mac, alias string) bool {
 	defer lan.Unlock()
 
 	mac = NormalizeMac(mac)
-	if e, found := lan.hosts[mac]; found {
 	lan.aliases.Set(mac, alias)
+	if e, found := lan.hosts[mac]; found {
 		e.Alias = alias
 		return true
 	}

network/wifi_ap.go

@@ -4,6 +4,8 @@ import (
 	"encoding/json"
 	"sync"
 	"time"
+
+	"github.com/evilsocket/islazy/data"
 )
 
 type AccessPoint struct {
@@ -65,7 +67,7 @@ func (ap *AccessPoint) RemoveClient(mac string) {
 	}
 }
 
-func (ap *AccessPoint) AddClientIfNew(bssid string, frequency int, rssi int8) (*Station, bool) {
+func (ap *AccessPoint) AddClientIfNew(bssid string, frequency int, rssi int8, aliases *data.UnsortedKV) (*Station, bool) {
 	ap.Lock()
 	defer ap.Unlock()
 
@@ -77,10 +79,17 @@ func (ap *AccessPoint) AddClientIfNew(bssid string, frequency int, rssi int8) (*
 		s.RSSI = rssi
 		s.LastSeen = time.Now()
 
+		if aliases != nil {
+			s.Alias = aliases.GetOr(bssid, "")
+		}
+
 		return s, false
 	}
 
 	s := NewStation("", bssid, frequency, rssi)
+	if aliases != nil {
+		s.Alias = aliases.GetOr(bssid, "")
+	}
 	ap.clients[bssid] = s
 
 	return s, true

packets/queue.go

@@ -1,9 +1,11 @@
 package packets
 
 import (
+	"encoding/json"
 	"fmt"
 	"net"
 	"sync"
+	"sync/atomic"
 
 	"github.com/bettercap/bettercap/network"
 
@@ -20,17 +22,15 @@ type Activity struct {
 }
 
 type Traffic struct {
-	Sent     uint64
+	Sent     uint64 `json:"sent"`
-	Received uint64
+	Received uint64 `json:"received"`
 }
 
 type Stats struct {
-	sync.RWMutex
+	Sent        uint64 `json:"sent"`
-
+	Received    uint64 `json:"received"`
-	Sent        uint64
+	PktReceived uint64 `json:"pkts_received"`
-	Received    uint64
+	Errors      uint64 `json:"errors"`
-	PktReceived uint64
-	Errors      uint64
 }
 
 type PacketCallback func(pkt gopacket.Packet)
@@ -38,11 +38,10 @@ type PacketCallback func(pkt gopacket.Packet)
 type Queue struct {
 	sync.RWMutex
 
-	Activities chan Activity `json:"-"`
+	Activities chan Activity
-
 	Stats      Stats
-	Protos  map[string]uint64
+	Protos     sync.Map
-	Traffic map[string]*Traffic
+	Traffic    sync.Map
 
 	iface      *network.Endpoint
 	handle     *pcap.Handle
@@ -53,10 +52,16 @@ type Queue struct {
 	active     bool
 }
 
+type queueJSON struct {
+	Stats   Stats               `json:"stats"`
+	Protos  map[string]int      `json:"protos"`
+	Traffic map[string]*Traffic `json:"traffic"`
+}
+
 func NewQueue(iface *network.Endpoint) (q *Queue, err error) {
 	q = &Queue{
-		Protos:     make(map[string]uint64),
+		Protos:     sync.Map{},
-		Traffic:    make(map[string]*Traffic),
+		Traffic:    sync.Map{},
 		Activities: make(chan Activity),
 
 		writes: &sync.WaitGroup{},
@@ -78,6 +83,28 @@ func NewQueue(iface *network.Endpoint) (q *Queue, err error) {
 	return
 }
 
+func (q *Queue) MarshalJSON() ([]byte, error) {
+	q.Lock()
+	defer q.Unlock()
+	doc := queueJSON{
+		Stats:   q.Stats,
+		Protos:  make(map[string]int),
+		Traffic: make(map[string]*Traffic),
+	}
+
+	q.Protos.Range(func(k, v interface{}) bool {
+		doc.Protos[k.(string)] = v.(int)
+		return true
+	})
+
+	q.Traffic.Range(func(k, v interface{}) bool {
+		doc.Traffic[k.(string)] = v.(*Traffic)
+		return true
+	})
+
+	return json.Marshal(doc)
+}
+
 func (q *Queue) OnPacket(cb PacketCallback) {
 	q.Lock()
 	defer q.Unlock()
@@ -102,14 +129,12 @@ func (q *Queue) trackProtocols(pkt gopacket.Packet) {
 			continue
 		}
 
-		q.Lock()
 		name := proto.String()
-		if _, found := q.Protos[name]; !found {
+		if v, found := q.Protos.Load(name); !found {
-			q.Protos[name] = 1
+			q.Protos.Store(name, 1)
 		} else {
-			q.Protos[name]++
+			q.Protos.Store(name, v.(int)+1)
 		}
-		q.Unlock()
 	}
 }
 
@@ -122,46 +147,34 @@ func (q *Queue) trackActivity(eth *layers.Ethernet, ip4 *layers.IPv4, address ne
 		Source: isSent,
 	}
 
-	q.Lock()
-	defer q.Unlock()
-
 	// initialize or update stats
 	addr := address.String()
-	if _, found := q.Traffic[addr]; !found {
+	if v, found := q.Traffic.Load(addr); !found {
 		if isSent {
-			q.Traffic[addr] = &Traffic{Sent: pktSize}
+			q.Traffic.Store(addr, &Traffic{Sent: pktSize})
 		} else {
-			q.Traffic[addr] = &Traffic{Received: pktSize}
+			q.Traffic.Store(addr, &Traffic{Received: pktSize})
 		}
 	} else {
 		if isSent {
-			q.Traffic[addr].Sent += pktSize
+			v.(*Traffic).Sent += pktSize
 		} else {
-			q.Traffic[addr].Received += pktSize
+			v.(*Traffic).Received += pktSize
 		}
 	}
 }
 
 func (q *Queue) TrackPacket(size uint64) {
-	q.Stats.Lock()
+	atomic.AddUint64(&q.Stats.PktReceived, 1)
-	defer q.Stats.Unlock()
+	atomic.AddUint64(&q.Stats.Received, size)
-
-	q.Stats.PktReceived++
-	q.Stats.Received += size
 }
 
 func (q *Queue) TrackSent(size uint64) {
-	q.Stats.Lock()
+	atomic.AddUint64(&q.Stats.Sent, size)
-	defer q.Stats.Unlock()
-
-	q.Stats.Sent += size
 }
 
 func (q *Queue) TrackError() {
-	q.Stats.Lock()
+	atomic.AddUint64(&q.Stats.Errors, 1)
-	defer q.Stats.Unlock()
-
-	q.Stats.Errors++
 }
 
 func (q *Queue) getPacketMeta(pkt gopacket.Packet) map[string]string {

session/session_core_handlers.go

@@ -274,12 +274,8 @@ func (s *Session) shHandler(args []string, sess *Session) error {
 func (s *Session) aliasHandler(args []string, sess *Session) error {
 	mac := args[0]
 	alias := str.Trim(args[1])
-
+	s.Lan.SetAliasFor(mac, alias)
-	if s.Lan.SetAliasFor(mac, alias) {
 	return nil
-	} else {
-		return fmt.Errorf("Could not find endpoint %s", mac)
-	}
 }
 
 func (s *Session) addHandler(h CommandHandler, c *readline.PrefixCompleter) {