From 58738b7723a23ed3e8d739d69b68f28dadc57ddf Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Thu, 24 Jan 2019 14:36:26 +0100
Subject: [PATCH] new: added wifi.show clients and encryption sorting

---
 modules/wifi.go           |  2 +-
 modules/wifi_show.go      |  5 ++-
 modules/wifi_show_sort.go | 72 +++++++++++++++++++++------------------
 3 files changed, 43 insertions(+), 36 deletions(-)

diff --git a/modules/wifi.go b/modules/wifi.go
index af044017..68c0cbbf 100644
--- a/modules/wifi.go
+++ b/modules/wifi.go
@@ -148,7 +148,7 @@ func NewWiFiModule(s *session.Session) *WiFiModule {
 		}))
 
 	w.selector = ViewSelectorFor(&w.SessionModule, "wifi.show",
-		[]string{"rssi", "bssid", "essid", "channel", "encryption", "seen", "sent", "rcvd"}, "rssi")
+		[]string{"rssi", "bssid", "essid", "channel", "encryption", "clients", "seen", "sent", "rcvd"}, "rssi")
 
 	w.AddHandler(session.NewModuleHandler("wifi.recon.channel", `wifi\.recon\.channel[\s]+([0-9]+(?:[, ]+[0-9]+)*|clear)`,
 		"WiFi channels (comma separated) or 'clear' for channel hopping.",
diff --git a/modules/wifi_show.go b/modules/wifi_show.go
index 67fe308e..b1edec7d 100644
--- a/modules/wifi_show.go
+++ b/modules/wifi_show.go
@@ -146,7 +146,6 @@ func (w *WiFiModule) doSelection() (err error, stations []*network.Station) {
 	}
 	stations = filtered
 
-	// "encryption"}, "rssi"
 	switch w.selector.SortBy {
 	case "seen":
 		sort.Sort(ByWiFiSeenSorter(stations))
@@ -156,6 +155,10 @@ func (w *WiFiModule) doSelection() (err error, stations []*network.Station) {
 		sort.Sort(ByBssidSorter(stations))
 	case "channel":
 		sort.Sort(ByChannelSorter(stations))
+	case "clients":
+		sort.Sort(ByClientsSorter(stations))
+	case "encryption":
+		sort.Sort(ByEncryptionSorter(stations))
 	case "sent":
 		sort.Sort(ByWiFiSentSorter(stations))
 	case "rcvd":
diff --git a/modules/wifi_show_sort.go b/modules/wifi_show_sort.go
index 0ba28c6c..5aa3bacf 100644
--- a/modules/wifi_show_sort.go
+++ b/modules/wifi_show_sort.go
@@ -2,7 +2,6 @@ package modules
 
 import (
 	"github.com/bettercap/bettercap/network"
-	"github.com/bettercap/bettercap/packets"
 	"github.com/bettercap/bettercap/session"
 )
 
@@ -11,7 +10,10 @@ type ByRSSISorter []*network.Station
 func (a ByRSSISorter) Len() int      { return len(a) }
 func (a ByRSSISorter) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
 func (a ByRSSISorter) Less(i, j int) bool {
-	return a[i].RSSI < a[j].RSSI
+	if a[i].RSSI == a[j].RSSI {
+		return a[i].HwAddress < a[j].HwAddress
+	}
+	return a[i].RSSI > a[j].RSSI
 }
 
 type ByChannelSorter []*network.Station
@@ -22,6 +24,17 @@ func (a ByChannelSorter) Less(i, j int) bool {
 	return a[i].Frequency < a[j].Frequency
 }
 
+type ByEncryptionSorter []*network.Station
+
+func (a ByEncryptionSorter) Len() int      { return len(a) }
+func (a ByEncryptionSorter) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
+func (a ByEncryptionSorter) Less(i, j int) bool {
+	if a[i].Encryption == a[j].Encryption {
+		return a[i].HwAddress < a[j].HwAddress
+	}
+	return a[i].Encryption < a[j].Encryption
+}
+
 type ByBssidSorter []*network.Station
 
 func (a ByBssidSorter) Len() int      { return len(a) }
@@ -54,22 +67,7 @@ type ByWiFiSentSorter []*network.Station
 func (a ByWiFiSentSorter) Len() int      { return len(a) }
 func (a ByWiFiSentSorter) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
 func (a ByWiFiSentSorter) Less(i, j int) bool {
-	session.I.Queue.Lock()
-	defer session.I.Queue.Unlock()
-
-	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
+	return a[i].Sent < a[j].Sent
 }
 
 type ByWiFiRcvdSorter []*network.Station
@@ -77,20 +75,26 @@ type ByWiFiRcvdSorter []*network.Station
 func (a ByWiFiRcvdSorter) Len() int      { return len(a) }
 func (a ByWiFiRcvdSorter) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
 func (a ByWiFiRcvdSorter) Less(i, j int) bool {
-	session.I.Queue.Lock()
-	defer session.I.Queue.Unlock()
-
-	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
+	return a[i].Received < a[j].Received
+}
+
+type ByClientsSorter []*network.Station
+
+func (a ByClientsSorter) Len() int      { return len(a) }
+func (a ByClientsSorter) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
+func (a ByClientsSorter) Less(i, j int) bool {
+	left := 0
+	right := 0
+
+	if ap, found := session.I.WiFi.Get(a[i].HwAddress); found {
+		left = ap.NumClients()
+	}
+	if ap, found := session.I.WiFi.Get(a[j].HwAddress); found {
+		right = ap.NumClients()
+	}
+
+	if left == right {
+		return a[i].HwAddress < a[j].HwAddress
+	}
+	return left < right
 }