From a4aa5acbcdd07d468589507853a947af0a2e9095 Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Mon, 18 Feb 2019 13:25:33 +0100
Subject: [PATCH] new: wifi.deauth and wifi.assoc now support autocompletion

---
 modules/ble/ble_recon.go | 15 ++------------
 modules/wifi/wifi.go     | 16 ++++++++++----
 session/session.go       | 45 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 17 deletions(-)

diff --git a/modules/ble/ble_recon.go b/modules/ble/ble_recon.go
index 8091a752..cd6c448a 100644
--- a/modules/ble/ble_recon.go
+++ b/modules/ble/ble_recon.go
@@ -8,7 +8,6 @@ import (
 	"fmt"
 	"io/ioutil"
 	golog "log"
-	"strings"
 	"time"
 
 	"github.com/bettercap/bettercap/modules/utils"
@@ -77,7 +76,7 @@ func NewBLERecon(s *session.Session) *BLERecon {
 			return mod.enumAllTheThings(network.NormalizeMac(args[0]))
 		})
 
-	enum.Complete("ble.enum", mod.macCompleter)
+	enum.Complete("ble.enum", s.BLECompleter)
 
 	mod.AddHandler(enum)
 
@@ -97,7 +96,7 @@ func NewBLERecon(s *session.Session) *BLERecon {
 			return mod.writeBuffer(mac, uuid, data)
 		})
 
-	write.Complete("ble.write", mod.macCompleter)
+	write.Complete("ble.write", s.BLECompleter)
 
 	mod.AddHandler(write)
 
@@ -116,16 +115,6 @@ func (mod BLERecon) Author() string {
 	return "Simone Margaritelli <evilsocket@gmail.com>"
 }
 
-func (mod *BLERecon) macCompleter(prefix string) []string {
-	macs := []string{""}
-	mod.Session.BLE.EachDevice(func(mac string, dev *network.BLEDevice) {
-		if prefix == "" || strings.HasPrefix(mac, prefix) {
-			macs = append(macs, mac)
-		}
-	})
-	return macs
-}
-
 func (mod *BLERecon) isEnumerating() bool {
 	return mod.currDevice != nil
 }
diff --git a/modules/wifi/wifi.go b/modules/wifi/wifi.go
index e790956c..92ce61e3 100644
--- a/modules/wifi/wifi.go
+++ b/modules/wifi/wifi.go
@@ -115,7 +115,7 @@ func NewWiFiModule(s *session.Session) *WiFiModule {
 		"-200",
 		"Minimum WiFi signal strength in dBm."))
 
-	mod.AddHandler(session.NewModuleHandler("wifi.deauth BSSID", `wifi\.deauth ((?:[a-fA-F0-9:]{11,})|all|\*)`,
+	deauth := session.NewModuleHandler("wifi.deauth BSSID", `wifi\.deauth ((?:[a-fA-F0-9:]{11,})|all|\*)`,
 		"Start a 802.11 deauth attack, if an access point BSSID is provided, every client will be deauthenticated, otherwise only the selected client. Use 'all', '*' or a broadcast BSSID (ff:ff:ff:ff:ff:ff) to iterate every access point with at least one client and start a deauth attack for each one.",
 		func(args []string) error {
 			if args[0] == "all" || args[0] == "*" {
@@ -126,7 +126,11 @@ func NewWiFiModule(s *session.Session) *WiFiModule {
 				return err
 			}
 			return mod.startDeauth(bssid)
-		}))
+		})
+
+	deauth.Complete("wifi.deauth", s.WiFiCompleterFull)
+
+	mod.AddHandler(deauth)
 
 	mod.AddParam(session.NewStringParameter("wifi.deauth.skip",
 		"",
@@ -141,7 +145,7 @@ func NewWiFiModule(s *session.Session) *WiFiModule {
 		"true",
 		"Send wifi deauth packets to open networks."))
 
-	mod.AddHandler(session.NewModuleHandler("wifi.assoc BSSID", `wifi\.assoc ((?:[a-fA-F0-9:]{11,})|all|\*)`,
+	assoc := session.NewModuleHandler("wifi.assoc BSSID", `wifi\.assoc ((?:[a-fA-F0-9:]{11,})|all|\*)`,
 		"Send an association request to the selected BSSID in order to receive a RSN PMKID key. Use 'all', '*' or a broadcast BSSID (ff:ff:ff:ff:ff:ff) to iterate for every access point.",
 		func(args []string) error {
 			if args[0] == "all" || args[0] == "*" {
@@ -152,7 +156,11 @@ func NewWiFiModule(s *session.Session) *WiFiModule {
 				return err
 			}
 			return mod.startAssoc(bssid)
-		}))
+		})
+
+	assoc.Complete("wifi.assoc", s.WiFiCompleter)
+
+	mod.AddHandler(assoc)
 
 	mod.AddParam(session.NewStringParameter("wifi.assoc.skip",
 		"",
diff --git a/session/session.go b/session/session.go
index 6f7544cb..79053290 100644
--- a/session/session.go
+++ b/session/session.go
@@ -154,6 +154,51 @@ func (s *Session) Unlock() {
 	s.WiFi.Unlock()
 }
 
+func (s *Session) LANCompleter(prefix string) []string {
+	macs := []string{""}
+	s.Lan.EachHost(func(mac string, e *network.Endpoint) {
+		if prefix == "" || strings.HasPrefix(mac, prefix) {
+			macs = append(macs, mac)
+		}
+	})
+	return macs
+}
+
+func (s *Session) WiFiCompleter(prefix string) []string {
+	macs := []string{""}
+	s.WiFi.EachAccessPoint(func(mac string, ap *network.AccessPoint) {
+		if prefix == "" || strings.HasPrefix(mac, prefix) {
+			macs = append(macs, mac)
+		}
+	})
+	return macs
+}
+
+func (s *Session) WiFiCompleterFull(prefix string) []string {
+	macs := []string{""}
+	s.WiFi.EachAccessPoint(func(mac string, ap *network.AccessPoint) {
+		if prefix == "" || strings.HasPrefix(mac, prefix) {
+			macs = append(macs, mac)
+		}
+		ap.EachClient(func(mac string, c *network.Station) {
+			if prefix == "" || strings.HasPrefix(mac, prefix) {
+				macs = append(macs, mac)
+			}
+		})
+	})
+	return macs
+}
+
+func (s *Session) BLECompleter(prefix string) []string {
+	macs := []string{""}
+	s.BLE.EachDevice(func(mac string, dev *network.BLEDevice) {
+		if prefix == "" || strings.HasPrefix(mac, prefix) {
+			macs = append(macs, mac)
+		}
+	})
+	return macs
+}
+
 func (s *Session) Module(name string) (err error, mod Module) {
 	for _, m := range s.Modules {
 		if m.Name() == name {