diff --git a/modules/ble/ble_recon.go b/modules/ble/ble_recon.go index f8422136..8570d7e3 100644 --- a/modules/ble/ble_recon.go +++ b/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 diff --git a/modules/can/can_recon.go b/modules/can/can_recon.go index 59ff34b8..c06a3d20 100644 --- a/modules/can/can_recon.go +++ b/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() diff --git a/modules/hid/hid.go b/modules/hid/hid.go index 08929303..e83cc091 100644 --- a/modules/hid/hid.go +++ b/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 { diff --git a/modules/hid/hid_recon.go b/modules/hid/hid_recon.go index 9473ca93..2d60d5a7 100644 --- a/modules/hid/hid_recon.go +++ b/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() diff --git a/modules/wifi/wifi.go b/modules/wifi/wifi.go index 45c6cddf..4b50bb86 100644 --- a/modules/wifi/wifi.go +++ b/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() diff --git a/session/module.go b/session/module.go index 343e6909..80daa120 100644 --- a/session/module.go +++ b/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) +} diff --git a/session/session.go b/session/session.go index 76c2df01..03c00636 100644 --- a/session/session.go +++ b/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 _, h := range m.Handlers() { - if parsed, args := h.Parse(line); parsed { - return h.Exec(args) + for _, mod := range s.Modules { + for _, modHandler := range mod.Handlers() { + if parsed, args := modHandler.Parse(line); parsed { + 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) }