new: ble, can, hid and wifi modules will now set a custom prompt (closes #1117)

modules/ble/ble_recon.go

@@ -81,7 +81,7 @@ func NewBLERecon(s *session.Session) *BLERecon {
 		"Enumerate services and characteristics for the given BLE device.",
 		func(args []string) error {
 			if mod.isEnumerating() {
-				return fmt.Errorf("An enumeration for %s is already running, please wait.", mod.currDevice.Device.ID())
+				return fmt.Errorf("an enumeration for %s is already running, please wait.", mod.currDevice.Device.ID())
 			}
 
 			mod.writeData = nil
@@ -100,11 +100,11 @@ func NewBLERecon(s *session.Session) *BLERecon {
 			mac := network.NormalizeMac(args[0])
 			uuid, err := gatt.ParseUUID(args[1])
 			if err != nil {
-				return fmt.Errorf("Error parsing %s: %s", args[1], err)
+				return fmt.Errorf("error parsing %s: %s", args[1], err)
 			}
 			data, err := hex.DecodeString(args[2])
 			if err != nil {
-				return fmt.Errorf("Error parsing %s: %s", args[2], err)
+				return fmt.Errorf("error parsing %s: %s", args[2], err)
 			}
 
 			return mod.writeBuffer(mac, uuid, data)
@@ -190,11 +190,15 @@ func (mod *BLERecon) Configure() (err error) {
 	return nil
 }
 
+const blePrompt = "{blb}{fw}BLE {fb}{reset} {bold}» {reset}"
+
 func (mod *BLERecon) Start() error {
 	if err := mod.Configure(); err != nil {
 		return err
 	}
 
+	mod.SetPrompt(blePrompt)
+
 	return mod.SetRunning(true, func() {
 		go mod.pruner()
 
@@ -221,6 +225,8 @@ func (mod *BLERecon) Start() error {
 }
 
 func (mod *BLERecon) Stop() error {
+	mod.SetPrompt(session.DefaultPrompt)
+
 	return mod.SetRunning(false, func() {
 		mod.quit <- true
 		<-mod.done

modules/can/can_recon.go

@@ -144,11 +144,15 @@ func (mod *CANModule) onFrame(frame can.Frame) {
 	mod.Session.Events.Add("can.message", msg)
 }
 
+const canPrompt = "{br}{fw}{env.can.device} {fb}{reset} {bold}» {reset}"
+
 func (mod *CANModule) Start() error {
 	if err := mod.Configure(); err != nil {
 		return err
 	}
 
+	mod.SetPrompt(canPrompt)
+
 	return mod.SetRunning(true, func() {
 		mod.Info("started on %s ...", mod.deviceName)
 
@@ -160,6 +164,8 @@ func (mod *CANModule) Start() error {
 }
 
 func (mod *CANModule) Stop() error {
+	mod.SetPrompt(session.DefaultPrompt)
+
 	return mod.SetRunning(false, func() {
 		if mod.conn != nil {
 			mod.recv.Close()

modules/hid/hid.go

@@ -237,6 +237,8 @@ func (mod *HIDRecon) forceStop() error {
 	})
 }
 func (mod *HIDRecon) Stop() error {
+	mod.SetPrompt(session.DefaultPrompt)
+
 	return mod.SetRunning(false, func() {
 		mod.waitGroup.Wait()
 		if mod.dongle != nil {

modules/hid/hid_recon.go

@@ -86,11 +86,15 @@ func (mod *HIDRecon) devPruner() {
 	}
 }
 
+const hidPrompt = "{by}{fw}HID {fb}{reset} {bold}» {reset}"
+
 func (mod *HIDRecon) Start() error {
 	if err := mod.Configure(); err != nil {
 		return err
 	}
 
+	mod.SetPrompt(hidPrompt)
+
 	return mod.SetRunning(true, func() {
 		mod.waitGroup.Add(1)
 		defer mod.waitGroup.Done()

modules/wifi/wifi.go

@@ -666,11 +666,15 @@ func (mod *WiFiModule) updateStats(dot11 *layers.Dot11, packet gopacket.Packet)
 	}
 }
 
+const wifiPrompt = "{by}{fb}{env.iface.name} {reset} {bold}» {reset}"
+
 func (mod *WiFiModule) Start() error {
 	if err := mod.Configure(); err != nil {
 		return err
 	}
 
+	mod.SetPrompt(wifiPrompt)
+
 	mod.SetRunning(true, func() {
 		// start channel hopper if needed
 		if mod.channel == 0 && mod.source == "" {
@@ -721,6 +725,8 @@ func (mod *WiFiModule) Start() error {
 }
 
 func (mod *WiFiModule) forcedStop() error {
+	mod.SetPrompt(session.DefaultPromptMonitor)
+
 	return mod.SetRunning(false, func() {
 		// signal the main for loop we want to exit
 		if !mod.pktSourceChanClosed {
@@ -732,6 +738,8 @@ func (mod *WiFiModule) forcedStop() error {
 }
 
 func (mod *WiFiModule) Stop() error {
+	mod.SetPrompt(session.DefaultPromptMonitor)
+
 	return mod.SetRunning(false, func() {
 		// wait any pending write operation
 		mod.writes.Wait()

session/module.go

@@ -17,6 +17,7 @@ type Module interface {
 	Name() string
 	Description() string
 	Author() string
+	Prompt() string
 	Handlers() []ModuleHandler
 	Parameters() map[string]*ModuleParam
 
@@ -29,33 +30,6 @@ type Module interface {
 
 type ModuleList []Module
 
-type moduleJSON struct {
-	Name        string                  `json:"name"`
-	Description string                  `json:"description"`
-	Author      string                  `json:"author"`
-	Parameters  map[string]*ModuleParam `json:"parameters"`
-	Handlers    []ModuleHandler         `json:"handlers"`
-	Running     bool                    `json:"running"`
-	State       map[string]interface{}  `json:"state"`
-}
-
-func (mm ModuleList) MarshalJSON() ([]byte, error) {
-	mods := []moduleJSON{}
-	for _, m := range mm {
-		mJSON := moduleJSON{
-			Name:        m.Name(),
-			Description: m.Description(),
-			Author:      m.Author(),
-			Parameters:  m.Parameters(),
-			Handlers:    m.Handlers(),
-			Running:     m.Running(),
-			State:       m.Extra(),
-		}
-		mods = append(mods, mJSON)
-	}
-	return json.Marshal(mods)
-}
-
 type SessionModule struct {
 	Name       string
 	Session    *Session
@@ -67,6 +41,7 @@ type SessionModule struct {
 	params   map[string]*ModuleParam
 	requires []string
 	tag      string
+	prompt   string
 }
 
 func AsTag(name string) string {
@@ -215,7 +190,7 @@ func (m SessionModule) DecParam(name string) (error, float64) {
 		}
 
 	} else {
-		return fmt.Errorf("Parameter %s does not exist.", name), 0
+		return fmt.Errorf("parameter %s does not exist", name), 0
 	}
 }
 
@@ -227,6 +202,14 @@ func (m SessionModule) BoolParam(name string) (error, bool) {
 	}
 }
 
+func (m *SessionModule) SetPrompt(prompt string) {
+	m.prompt = prompt
+}
+
+func (m *SessionModule) Prompt() string {
+	return m.prompt
+}
+
 func (m *SessionModule) AddHandler(h ModuleHandler) {
 	m.handlers = append(m.handlers, h)
 }
@@ -302,3 +285,30 @@ func (m *SessionModule) SetRunning(running bool, cb func()) error {
 
 	return nil
 }
+
+type moduleJSON struct {
+	Name        string                  `json:"name"`
+	Description string                  `json:"description"`
+	Author      string                  `json:"author"`
+	Parameters  map[string]*ModuleParam `json:"parameters"`
+	Handlers    []ModuleHandler         `json:"handlers"`
+	Running     bool                    `json:"running"`
+	State       map[string]interface{}  `json:"state"`
+}
+
+func (mm ModuleList) MarshalJSON() ([]byte, error) {
+	mods := []moduleJSON{}
+	for _, m := range mm {
+		mJSON := moduleJSON{
+			Name:        m.Name(),
+			Description: m.Description(),
+			Author:      m.Author(),
+			Parameters:  m.Parameters(),
+			Handlers:    m.Handlers(),
+			Running:     m.Running(),
+			State:       m.Extra(),
+		}
+		mods = append(mods, mJSON)
+	}
+	return json.Marshal(mods)
+}

session/session.go

@@ -458,10 +458,18 @@ func (s *Session) Run(line string) error {
 	}
 
 	// is it a module command?
-	for _, m := range s.Modules {
+	for _, mod := range s.Modules {
-		for _, h := range m.Handlers() {
+		for _, modHandler := range mod.Handlers() {
-			if parsed, args := h.Parse(line); parsed {
+			if parsed, args := modHandler.Parse(line); parsed {
-				return h.Exec(args)
+				if err := modHandler.Exec(args); err != nil {
+					return err
+				} else if prompt := mod.Prompt(); prompt != "" {
+					// if the module handler has been executed successfully and
+					// the module overrides the prompt, set it
+					s.Env.Set(PromptVariable, prompt)
+					s.Refresh()
+				}
+				return nil
 			}
 		}
 	}
@@ -478,5 +486,5 @@ func (s *Session) Run(line string) error {
 		return nil
 	}
 
-	return fmt.Errorf("unknown or invalid syntax \"%s%s%s\", type %shelp%s for the help menu.", tui.BOLD, line, tui.RESET, tui.BOLD, tui.RESET)
+	return fmt.Errorf("unknown or invalid syntax \"%s%s%s\", type %shelp%s for the help menu", tui.BOLD, line, tui.RESET, tui.BOLD, tui.RESET)
 }