bettercap/session/module.go

package session

import (
	"encoding/json"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	"github.com/evilsocket/islazy/log"
	"github.com/evilsocket/islazy/str"
	"github.com/evilsocket/islazy/tui"
)

type Module interface {
	Name() string
	Description() string
	Author() string
	Prompt() string
	Handlers() []ModuleHandler
	Parameters() map[string]*ModuleParam

	Extra() map[string]interface{}
	Required() []string
	Running() bool
	Start() error
	Stop() error
}

type ModuleList []Module

type SessionModule struct {
	Name       string
	Session    *Session
	Started    bool
	StatusLock *sync.RWMutex
	State      *sync.Map

	handlers []ModuleHandler
	params   map[string]*ModuleParam
	requires []string
	tag      string
	prompt   string
}

func AsTag(name string) string {
	return fmt.Sprintf("%s ", tui.Wrap(tui.BACKLIGHTBLUE, tui.Wrap(tui.FOREBLACK, name)))
}

func NewSessionModule(name string, s *Session) SessionModule {
	m := SessionModule{
		Name:       name,
		Session:    s,
		Started:    false,
		StatusLock: &sync.RWMutex{},
		State:      &sync.Map{},

		requires: make([]string, 0),
		handlers: make([]ModuleHandler, 0),
		params:   make(map[string]*ModuleParam),
		tag:      AsTag(name),
	}

	return m
}

func (m *SessionModule) Extra() map[string]interface{} {
	extra := make(map[string]interface{})
	m.State.Range(func(k, v interface{}) bool {
		extra[k.(string)] = v
		return true
	})
	return extra
}

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) Debug(format string, args ...interface{}) {
	m.Session.Events.Log(log.DEBUG, m.tag+format, args...)
}

func (m *SessionModule) Info(format string, args ...interface{}) {
	m.Session.Events.Log(log.INFO, m.tag+format, args...)
}

func (m *SessionModule) Warning(format string, args ...interface{}) {
	m.Session.Events.Log(log.WARNING, m.tag+format, args...)
}

func (m *SessionModule) Error(format string, args ...interface{}) {
	m.Session.Events.Log(log.ERROR, m.tag+format, args...)
}

func (m *SessionModule) Fatal(format string, args ...interface{}) {
	m.Session.Events.Log(log.FATAL, m.tag+format, args...)
}

func (m *SessionModule) Printf(format string, a ...interface{}) {
	m.Session.Events.Printf(format, a...)
}

func (m *SessionModule) Requires(modName string) {
	m.requires = append(m.requires, modName)
}

func (m *SessionModule) Required() []string {
	return m.requires
}

func (m *SessionModule) Handlers() []ModuleHandler {
	return m.handlers
}

func (m *SessionModule) Parameters() map[string]*ModuleParam {
	return m.params
}

func (m *SessionModule) Param(name string) *ModuleParam {
	return m.params[name]
}

func (m SessionModule) ListParam(name string) (err error, values []string) {
	values = make([]string, 0)
	list := ""
	if err, list = m.StringParam(name); err != nil {
		return
	} else {
		parts := strings.Split(list, ",")
		for _, part := range parts {
			part = str.Trim(part)
			if part != "" {
				values = append(values, part)
			}
		}
	}
	return
}

func (m SessionModule) StringParam(name string) (error, string) {
	if p, found := m.params[name]; found {
		if err, v := p.Get(m.Session); err != nil {
			return err, ""
		} else {
			return nil, v.(string)
		}
	} else {
		return fmt.Errorf("Parameter %s does not exist.", name), ""
	}
}

func (m SessionModule) IPParam(name string) (error, net.IP) {
	if err, v := m.StringParam(name); err != nil {
		return err, nil
	} else {
		return nil, net.ParseIP(v)
	}
}

func (m SessionModule) IntParam(name string) (error, int) {
	if p, found := m.params[name]; found {
		if err, v := p.Get(m.Session); err != nil {
			return err, 0
		} else {
			return nil, v.(int)
		}

	} else {
		return fmt.Errorf("Parameter %s does not exist.", name), 0
	}
}

func (m SessionModule) DecParam(name string) (error, float64) {
	if p, found := m.params[name]; found {
		if err, v := p.Get(m.Session); err != nil {
			return err, 0
		} else {
			return nil, v.(float64)
		}

	} else {
		return fmt.Errorf("parameter %s does not exist", name), 0
	}
}

func (m SessionModule) BoolParam(name string) (error, bool) {
	if err, v := m.params[name].Get(m.Session); err != nil {
		return err, false
	} else {
		return nil, v.(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)
}

func (m *SessionModule) AddParam(p *ModuleParam) *ModuleParam {
	m.params[p.Name] = p
	p.Register(m.Session)
	return p
}

func (m *SessionModule) AddObservableParam(p *ModuleParam, cb EnvironmentChangedCallback) *ModuleParam {
	m.params[p.Name] = p
	p.RegisterObserver(m.Session, cb)
	return p
}

func (m *SessionModule) Running() bool {
	m.StatusLock.RLock()
	defer m.StatusLock.RUnlock()
	return m.Started
}

func (m *SessionModule) SetRunning(running bool, cb func()) error {
	if running == m.Running() {
		if m.Started {
			return ErrAlreadyStarted(m.Name)
		} else {
			return ErrAlreadyStopped(m.Name)
		}
	}

	if running == true {
		for _, modName := range m.Required() {
			if m.Session.IsOn(modName) == false {
				m.Info("starting %s as a requirement for %s", modName, m.Name)
				if err := m.Session.Run(modName + " on"); err != nil {
					return fmt.Errorf("error while starting module %s as a requirement for %s: %v", modName, m.Name, err)
				}
			}
		}
	}

	m.StatusLock.Lock()
	m.Started = running
	m.StatusLock.Unlock()

	if running {
		m.Session.Events.Add("mod.started", m.Name)
	} else {
		m.Session.Events.Add("mod.stopped", m.Name)
	}

	if cb != nil {
		if running {
			// this is the worker, start async
			go cb()
		} else {
			// stop callback, this is sync with a 10 seconds timeout
			done := make(chan bool, 1)
			go func() {
				cb()
				done <- true
			}()

			select {
			case <-done:
				return nil
			case <-time.After(10 * time.Second):
				fmt.Printf("%s: Stopping module %s timed out.\n", tui.Yellow(tui.Bold("WARNING")), m.Name)
			}
		}
	}

	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)
}