From 756c04fd95af1ab1a254b1078d47b7a7142b07fc Mon Sep 17 00:00:00 2001
From: evilsocket <evilsocket@gmail.com>
Date: Mon, 18 Mar 2019 12:07:00 +0100
Subject: [PATCH] new: module can now export a State map with specific
 information

---
 modules/ble/ble_recon.go |  4 ++++
 modules/wifi/wifi.go     | 23 +++++++++++++++++----
 session/module.go        | 44 ++++++++++++++++++++++++++++++++++++----
 3 files changed, 63 insertions(+), 8 deletions(-)

diff --git a/modules/ble/ble_recon.go b/modules/ble/ble_recon.go
index 1bea0d99..d6a5cddd 100644
--- a/modules/ble/ble_recon.go
+++ b/modules/ble/ble_recon.go
@@ -42,6 +42,8 @@ func NewBLERecon(s *session.Session) *BLERecon {
 		connected:     false,
 	}
 
+	mod.InitState("scanning")
+
 	mod.selector = utils.ViewSelectorFor(&mod.SessionModule,
 		"ble.show",
 		[]string{"rssi", "mac", "seen"}, "rssi asc")
@@ -196,6 +198,7 @@ func (mod *BLERecon) Stop() error {
 		mod.Debug("module stopped, cleaning state")
 		mod.gattDevice = nil
 		mod.setCurrentDevice(nil)
+		mod.ResetState()
 	})
 }
 
@@ -216,6 +219,7 @@ func (mod *BLERecon) pruner() {
 func (mod *BLERecon) setCurrentDevice(dev *network.BLEDevice) {
 	mod.connected = false
 	mod.currDevice = dev
+	mod.State.Store("scanning", dev)
 }
 
 func (mod *BLERecon) writeBuffer(mac string, uuid gatt.UUID, data []byte) error {
diff --git a/modules/wifi/wifi.go b/modules/wifi/wifi.go
index bdd513d5..c5d5ab90 100644
--- a/modules/wifi/wifi.go
+++ b/modules/wifi/wifi.go
@@ -81,6 +81,8 @@ func NewWiFiModule(s *session.Session) *WiFiModule {
 		chanLock:      &sync.Mutex{},
 	}
 
+	mod.InitState("channels")
+
 	mod.AddParam(session.NewStringParameter("wifi.interface",
 		"",
 		"",
@@ -124,7 +126,8 @@ func NewWiFiModule(s *session.Session) *WiFiModule {
 		func(args []string) (err error) {
 			mod.ap = nil
 			mod.stickChan = 0
-			mod.frequencies, err = network.GetSupportedFrequencies(mod.iface.Name())
+			freqs, err := network.GetSupportedFrequencies(mod.iface.Name())
+			mod.setFrequencies(freqs)
 			mod.hopChanges <- true
 			return err
 		}))
@@ -285,8 +288,7 @@ func NewWiFiModule(s *session.Session) *WiFiModule {
 				}
 			}
 
-			mod.Debug("new frequencies: %v", freqs)
-			mod.frequencies = freqs
+			mod.setFrequencies(freqs)
 
 			// if wifi.recon is not running, this would block forever
 			if mod.Running() {
@@ -330,6 +332,17 @@ const (
 	ErrIfaceNotUp = "Interface Not Up"
 )
 
+func (mod *WiFiModule) setFrequencies(freqs []int) {
+	mod.Debug("new frequencies: %v", freqs)
+
+	mod.frequencies = freqs
+	channels := []int{}
+	for _, freq := range freqs {
+		channels = append(channels, network.Dot11Freq2Chan(freq))
+	}
+	mod.State.Store("channels", channels)
+}
+
 func (mod *WiFiModule) Configure() error {
 	var ifName string
 	var hopPeriod int
@@ -430,8 +443,10 @@ func (mod *WiFiModule) Configure() error {
 	if mod.source == "" {
 		// No channels setted, retrieve frequencies supported by the card
 		if len(mod.frequencies) == 0 {
-			if mod.frequencies, err = network.GetSupportedFrequencies(ifName); err != nil {
+			if freqs, err := network.GetSupportedFrequencies(ifName); err != nil {
 				return fmt.Errorf("error while getting supported frequencies of %s: %s", ifName, err)
+			} else {
+				mod.setFrequencies(freqs)
 			}
 
 			mod.Debug("wifi supported frequencies: %v", mod.frequencies)
diff --git a/session/module.go b/session/module.go
index 5accba3d..f8433551 100644
--- a/session/module.go
+++ b/session/module.go
@@ -1,6 +1,7 @@
 package session
 
 import (
+	"encoding/json"
 	"fmt"
 	"net"
 	"strings"
@@ -25,16 +26,23 @@ type Module interface {
 }
 
 type SessionModule struct {
-	Name       string        `json:"name"`
-	Session    *Session      `json:"-"`
-	Started    bool          `json:"started"`
-	StatusLock *sync.RWMutex `json:"-"`
+	Name       string
+	Session    *Session
+	Started    bool
+	StatusLock *sync.RWMutex
+	State      sync.Map
 
 	handlers []ModuleHandler
 	params   map[string]*ModuleParam
 	tag      string
 }
 
+type sessionModuleJSON struct {
+	Name    string                 `json:"name"`
+	Started bool                   `json:"started"`
+	State   map[string]interface{} `json:"state"`
+}
+
 func AsTag(name string) string {
 	return fmt.Sprintf("%s ", tui.Wrap(tui.BACKLIGHTBLUE, tui.Wrap(tui.FOREBLACK, name)))
 }
@@ -54,6 +62,34 @@ func NewSessionModule(name string, s *Session) SessionModule {
 	return m
 }
 
+func (m *SessionModule) InitState(keys ...string) {
+	for _, key := range keys {
+		m.State.Store(key, nil)
+	}
+}
+
+func (m *SessionModule) ResetState() {
+	m.State.Range(func(k, v interface{}) bool {
+		m.State.Store(k, nil)
+		return true
+	})
+}
+
+func (m *SessionModule) MarshalJSON() ([]byte, error) {
+	doc := sessionModuleJSON{
+		Name:    m.Name,
+		Started: m.Started,
+		State:   make(map[string]interface{}),
+	}
+
+	m.State.Range(func(k, v interface{}) bool {
+		doc.State[k.(string)] = v
+		return true
+	})
+
+	return json.Marshal(doc)
+}
+
 func (m *SessionModule) Debug(format string, args ...interface{}) {
 	m.Session.Events.Log(log.DEBUG, m.tag+format, args...)
 }