From 49e2116d4644da3333f749d14157741fd0b6158b Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Mon, 18 Feb 2019 13:06:41 +0100
Subject: [PATCH] new: ble.enum and ble.write now support autocompletion

---
 modules/ble/ble_recon.go  | 27 +++++++++++++++++++++++----
 network/ble.go            |  9 +++++++++
 session/module_handler.go | 11 +++++++++++
 session/session_setup.go  | 20 ++++++++++++--------
 4 files changed, 55 insertions(+), 12 deletions(-)

diff --git a/modules/ble/ble_recon.go b/modules/ble/ble_recon.go
index 650e087d..8091a752 100644
--- a/modules/ble/ble_recon.go
+++ b/modules/ble/ble_recon.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	golog "log"
+	"strings"
 	"time"
 
 	"github.com/bettercap/bettercap/modules/utils"
@@ -63,7 +64,7 @@ func NewBLERecon(s *session.Session) *BLERecon {
 			return mod.Show()
 		}))
 
-	mod.AddHandler(session.NewModuleHandler("ble.enum MAC", "ble.enum "+network.BLEMacValidator,
+	enum := session.NewModuleHandler("ble.enum MAC", "ble.enum "+network.BLEMacValidator,
 		"Enumerate services and characteristics for the given BLE device.",
 		func(args []string) error {
 			if mod.isEnumerating() {
@@ -74,9 +75,13 @@ func NewBLERecon(s *session.Session) *BLERecon {
 			mod.writeUUID = nil
 
 			return mod.enumAllTheThings(network.NormalizeMac(args[0]))
-		}))
+		})
 
-	mod.AddHandler(session.NewModuleHandler("ble.write MAC UUID HEX_DATA", "ble.write "+network.BLEMacValidator+" ([a-fA-F0-9]+) ([a-fA-F0-9]+)",
+	enum.Complete("ble.enum", mod.macCompleter)
+
+	mod.AddHandler(enum)
+
+	write := session.NewModuleHandler("ble.write MAC UUID HEX_DATA", "ble.write "+network.BLEMacValidator+" ([a-fA-F0-9]+) ([a-fA-F0-9]+)",
 		"Write the HEX_DATA buffer to the BLE device with the specified MAC address, to the characteristics with the given UUID.",
 		func(args []string) error {
 			mac := network.NormalizeMac(args[0])
@@ -90,7 +95,11 @@ func NewBLERecon(s *session.Session) *BLERecon {
 			}
 
 			return mod.writeBuffer(mac, uuid, data)
-		}))
+		})
+
+	write.Complete("ble.write", mod.macCompleter)
+
+	mod.AddHandler(write)
 
 	return mod
 }
@@ -107,6 +116,16 @@ 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/network/ble.go b/network/ble.go
index aee08e1a..6f1592c9 100644
--- a/network/ble.go
+++ b/network/ble.go
@@ -95,3 +95,12 @@ func (b *BLE) Devices() (devices []*BLEDevice) {
 	}
 	return
 }
+
+func (b *BLE) EachDevice(cb func(mac string, d *BLEDevice)) {
+	b.Lock()
+	defer b.Unlock()
+
+	for m, dev := range b.devices {
+		cb(m, dev)
+	}
+}
diff --git a/session/module_handler.go b/session/module_handler.go
index 422963a6..063f98dc 100644
--- a/session/module_handler.go
+++ b/session/module_handler.go
@@ -6,7 +6,10 @@ import (
 	"regexp"
 	"strconv"
 
+	"github.com/evilsocket/islazy/str"
 	"github.com/evilsocket/islazy/tui"
+
+	"github.com/bettercap/readline"
 )
 
 const IPv4Validator = `^(?:[0-9]{1,3}\.){3}[0-9]{1,3}$`
@@ -16,6 +19,7 @@ type ModuleHandler struct {
 	Description string
 	Parser      *regexp.Regexp
 	Exec        func(args []string) error
+	Completer   *readline.PrefixCompleter
 }
 
 func NewModuleHandler(name string, expr string, desc string, exec func(args []string) error) ModuleHandler {
@@ -33,6 +37,13 @@ func NewModuleHandler(name string, expr string, desc string, exec func(args []st
 	return h
 }
 
+func (h *ModuleHandler) Complete(name string, cb func(prefix string) []string) {
+	h.Completer = readline.PcItem(name, readline.PcItemDynamic(func(prefix string) []string {
+		prefix = str.Trim(prefix[len(name):])
+		return cb(prefix)
+	}))
+}
+
 func (h *ModuleHandler) Help(padding int) string {
 	return fmt.Sprintf("  "+tui.Bold("%"+strconv.Itoa(padding)+"s")+" : %s\n", h.Name, h.Description)
 }
diff --git a/session/session_setup.go b/session/session_setup.go
index d17fb267..4d5b6308 100644
--- a/session/session_setup.go
+++ b/session/session_setup.go
@@ -38,17 +38,21 @@ func (s *Session) setupReadline() (err error) {
 	tree := make(map[string][]string)
 	for _, m := range s.Modules {
 		for _, h := range m.Handlers() {
-			parts := strings.Split(h.Name, " ")
-			name := parts[0]
+			if h.Completer == nil {
+				parts := strings.Split(h.Name, " ")
+				name := parts[0]
 
-			if _, found := tree[name]; !found {
-				tree[name] = []string{}
-			}
+				if _, found := tree[name]; !found {
+					tree[name] = []string{}
+				}
 
-			var appendedOption = strings.Join(parts[1:], " ")
+				var appendedOption = strings.Join(parts[1:], " ")
 
-			if len(appendedOption) > 0 && !containsCapitals(appendedOption) {
-				tree[name] = append(tree[name], appendedOption)
+				if len(appendedOption) > 0 && !containsCapitals(appendedOption) {
+					tree[name] = append(tree[name], appendedOption)
+				}
+			} else {
+				prefixCompleters = append(prefixCompleters, h.Completer)
 			}
 		}
 	}