Init

Signed-off-by: Hadi <112569860+anotherhadi@users.noreply.github.com>

internal/plugins/cmd.go

@@ -0,0 +1,15 @@
+package plugins
+
+import tea "charm.land/bubbletea/v2"
+
+func WaitForNotif(mgr *Manager) tea.Cmd {
+	return func() tea.Msg {
+		return <-mgr.Notifs
+	}
+}
+
+func WaitForQuit(mgr *Manager) tea.Cmd {
+	return func() tea.Msg {
+		return PluginQuitMsg{Reason: <-mgr.Quit}
+	}
+}

internal/plugins/lua.go

@@ -0,0 +1,206 @@
+package plugins
+
+import (
+	"log"
+	"net/url"
+	"strings"
+	"time"
+
+	"github.com/anotherhadi/spilltea/internal/db"
+	goproxy "github.com/lqqyt2423/go-mitmproxy/proxy"
+	lua "github.com/yuin/gopher-lua"
+)
+
+func newLuaState(mgr *Manager, p *Plugin) *lua.LState {
+	L := lua.NewState()
+	registerUtilities(L, mgr, p)
+	return L
+}
+
+func registerUtilities(L *lua.LState, mgr *Manager, p *Plugin) {
+	L.SetGlobal("log", L.NewFunction(func(L *lua.LState) int {
+		msg := L.CheckString(1)
+		log.Printf("[plugin:%s] %s", p.Name, msg)
+		return 0
+	}))
+
+	L.SetGlobal("notif", L.NewFunction(func(L *lua.LState) int {
+		title := L.CheckString(1)
+		body := L.CheckString(2)
+		select {
+		case mgr.Notifs <- PluginNotifMsg{Title: title, Body: body}:
+		default:
+		}
+		return 0
+	}))
+
+	L.SetGlobal("create_finding", L.NewFunction(func(L *lua.LState) int {
+		t := L.CheckTable(1)
+		title := luaTableString(t, "title")
+		desc := luaTableString(t, "description")
+		key := luaTableString(t, "key")
+		severity := luaTableString(t, "severity")
+		if severity == "" {
+			severity = "info"
+		}
+		if key == "" {
+			key = title
+		}
+		if mgr.db == nil {
+			return 0
+		}
+		inserted, err := mgr.db.UpsertFinding(db.Finding{
+			PluginName:  p.Name,
+			DedupKey:    key,
+			Title:       title,
+			Description: desc,
+			Severity:    severity,
+			CreatedAt:   time.Now(),
+		})
+		if err != nil {
+			log.Printf("[plugin:%s] create_finding error: %v", p.Name, err)
+			return 0
+		}
+		_ = inserted
+		return 0
+	}))
+
+	L.SetGlobal("is_in_scope", L.NewFunction(func(L *lua.LState) int {
+		raw := L.CheckString(1)
+		if mgr.broker == nil {
+			L.Push(lua.LTrue)
+			return 1
+		}
+		u, err := url.Parse(raw)
+		if err != nil {
+			L.Push(lua.LFalse)
+			return 1
+		}
+		path := u.Path
+		if path == "" {
+			path = "/"
+		}
+		L.Push(lua.LBool(mgr.broker.IsInScope(u.Host + path)))
+		return 1
+	}))
+
+	L.SetGlobal("quit", L.NewFunction(func(L *lua.LState) int {
+		reason := L.OptString(1, "plugin requested quit")
+		select {
+		case mgr.Quit <- reason:
+		default:
+		}
+		return 0
+	}))
+}
+
+func luaTableString(t *lua.LTable, key string) string {
+	v := t.RawGetString(key)
+	if s, ok := v.(lua.LString); ok {
+		return string(s)
+	}
+	return ""
+}
+
+func pushRequest(L *lua.LState, f *goproxy.Flow) *lua.LTable {
+	t := L.NewTable()
+	r := f.Request
+	L.SetField(t, "method", lua.LString(r.Method))
+	L.SetField(t, "url", lua.LString(r.URL.String()))
+	L.SetField(t, "host", lua.LString(r.URL.Host))
+	L.SetField(t, "path", lua.LString(r.URL.Path))
+
+	headers := L.NewTable()
+	for k, vals := range r.Header {
+		L.SetField(headers, k, lua.LString(strings.Join(vals, ", ")))
+	}
+	L.SetField(t, "headers", headers)
+
+	L.SetField(t, "get_body", L.NewFunction(func(L *lua.LState) int {
+		L.Push(lua.LString(string(r.Body)))
+		return 1
+	}))
+
+	L.SetField(t, "set_header", L.NewFunction(func(L *lua.LState) int {
+		name := L.CheckString(2)
+		value := L.CheckString(3)
+		r.Header.Set(name, value)
+		return 0
+	}))
+
+	L.SetField(t, "set_body", L.NewFunction(func(L *lua.LState) int {
+		body := L.CheckString(2)
+		r.Body = []byte(body)
+		return 0
+	}))
+
+	return t
+}
+
+func pushResponse(L *lua.LState, f *goproxy.Flow) *lua.LTable {
+	t := L.NewTable()
+	if f.Response == nil {
+		return t
+	}
+	resp := f.Response
+	L.SetField(t, "status_code", lua.LNumber(resp.StatusCode))
+
+	headers := L.NewTable()
+	for k, vals := range resp.Header {
+		L.SetField(headers, k, lua.LString(strings.Join(vals, ", ")))
+	}
+	L.SetField(t, "headers", headers)
+
+	L.SetField(t, "get_body", L.NewFunction(func(L *lua.LState) int {
+		L.Push(lua.LString(string(resp.Body)))
+		return 1
+	}))
+
+	L.SetField(t, "set_header", L.NewFunction(func(L *lua.LState) int {
+		name := L.CheckString(2)
+		value := L.CheckString(3)
+		resp.Header.Set(name, value)
+		return 0
+	}))
+
+	L.SetField(t, "set_body", L.NewFunction(func(L *lua.LState) int {
+		body := L.CheckString(2)
+		resp.Body = []byte(body)
+		return 0
+	}))
+
+	return t
+}
+
+func pushEntry(L *lua.LState, e db.Entry) *lua.LTable {
+	t := L.NewTable()
+	L.SetField(t, "id", lua.LNumber(e.ID))
+	L.SetField(t, "method", lua.LString(e.Method))
+	L.SetField(t, "host", lua.LString(e.Host))
+	L.SetField(t, "path", lua.LString(e.Path))
+	L.SetField(t, "status_code", lua.LNumber(e.StatusCode))
+	L.SetField(t, "timestamp", lua.LString(e.Timestamp.Format("2006-01-02 15:04:05")))
+	L.SetField(t, "request_raw", lua.LString(e.RequestRaw))
+	L.SetField(t, "response_raw", lua.LString(e.ResponseRaw))
+	return t
+}
+
+func callHook(p *Plugin, hookName string, args ...lua.LValue) (string, error) {
+	fn := p.L.GetGlobal(hookName)
+	if fn == lua.LNil {
+		return "", nil
+	}
+	if err := p.L.CallByParam(lua.P{
+		Fn:      fn,
+		NRet:    1,
+		Protect: true,
+	}, args...); err != nil {
+		return "", err
+	}
+	ret := p.L.Get(-1)
+	p.L.Pop(1)
+	if s, ok := ret.(lua.LString); ok {
+		return string(s), nil
+	}
+	return "", nil
+}

internal/plugins/manager.go

@@ -0,0 +1,346 @@
+package plugins
+
+import (
+	"fmt"
+	"log"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+
+	"github.com/anotherhadi/spilltea/internal/db"
+	"github.com/anotherhadi/spilltea/internal/intercept"
+	goproxy "github.com/lqqyt2423/go-mitmproxy/proxy"
+	lua "github.com/yuin/gopher-lua"
+)
+
+type Manager struct {
+	mu      sync.RWMutex
+	plugins []*Plugin
+
+	db     *db.DB
+	broker *intercept.Broker
+
+	Notifs chan PluginNotifMsg
+	Quit   chan string
+}
+
+func NewManager(broker *intercept.Broker) *Manager {
+	mgr := &Manager{
+		broker:   broker,
+		Notifs: make(chan PluginNotifMsg, 64),
+		Quit:   make(chan string, 4),
+	}
+	if broker != nil {
+		broker.SetOnNewEntry(mgr.RunOnHistoryEntry)
+	}
+	return mgr
+}
+
+func (m *Manager) SetDB(d *db.DB) {
+	m.db = d
+}
+
+func (m *Manager) LoadFromDir(dir string) error {
+	entries, err := os.ReadDir(dir)
+	if os.IsNotExist(err) {
+		return nil
+	}
+	if err != nil {
+		return err
+	}
+
+	var states map[string]db.PluginState
+	if m.db != nil {
+		list, err := m.db.LoadPluginStates()
+		if err == nil {
+			states = make(map[string]db.PluginState, len(list))
+			for _, s := range list {
+				states[s.Name] = s
+			}
+		}
+	}
+
+	for _, e := range entries {
+		if e.IsDir() || !strings.HasSuffix(e.Name(), ".lua") {
+			continue
+		}
+		path := filepath.Join(dir, e.Name())
+		p, err := m.loadPlugin(path)
+		if err != nil {
+			log.Printf("plugin load error %s: %v", path, err)
+			continue
+		}
+		if s, ok := states[p.Name]; ok {
+			p.Enabled = s.Enabled
+			p.ConfigText = s.ConfigText
+		}
+		m.mu.Lock()
+		m.plugins = append(m.plugins, p)
+		m.mu.Unlock()
+	}
+	return nil
+}
+
+func (m *Manager) loadPlugin(path string) (*Plugin, error) {
+	p := &Plugin{
+		FilePath: path,
+		Enabled:  true,
+		hooks:    make(map[string]HookConfig),
+	}
+	p.L = newLuaState(m, p)
+	if err := p.L.DoFile(path); err != nil {
+		p.L.Close()
+		return nil, err
+	}
+
+	pluginTable, ok := p.L.GetGlobal("Plugin").(*lua.LTable)
+	if !ok {
+		p.L.Close()
+		return nil, fmt.Errorf("missing Plugin table")
+	}
+
+	if s, ok := pluginTable.RawGetString("name").(lua.LString); ok {
+		p.Name = string(s)
+	}
+	if p.Name == "" {
+		p.Name = strings.TrimSuffix(filepath.Base(path), ".lua")
+	}
+
+	// Defaults when not overridden by the Plugin table.
+	hookDefaults := map[string]bool{
+		"on_start":         true,  // always sync
+		"on_request":       false, // async
+		"on_response":      false, // async
+		"on_quit":          true,  // always sync
+		"on_history_entry": false, // always async
+	}
+	for hookName, defaultSync := range hookDefaults {
+		// Plugin table entry overrides the default (except on_start/on_quit/on_history_entry which are fixed).
+		if hookName != "on_start" && hookName != "on_quit" && hookName != "on_history_entry" {
+			if tbl, ok := pluginTable.RawGetString(hookName).(*lua.LTable); ok {
+				p.hooks[hookName] = HookConfig{Sync: tbl.RawGetString("sync") == lua.LTrue}
+				continue
+			}
+		}
+		// Auto-detect: register the hook if the function exists as a global.
+		if p.L.GetGlobal(hookName) != lua.LNil {
+			p.hooks[hookName] = HookConfig{Sync: defaultSync}
+		}
+	}
+
+	return p, nil
+}
+
+func (m *Manager) GetPlugins() []*Plugin {
+	m.mu.RLock()
+	defer m.mu.RUnlock()
+	out := make([]*Plugin, len(m.plugins))
+	copy(out, m.plugins)
+	return out
+}
+
+func (m *Manager) TogglePlugin(name string) {
+	m.mu.RLock()
+	var found *Plugin
+	for _, p := range m.plugins {
+		if p.Name == name {
+			found = p
+			break
+		}
+	}
+	m.mu.RUnlock()
+	if found == nil {
+		return
+	}
+	found.mu.Lock()
+	found.Enabled = !found.Enabled
+	enabled := found.Enabled
+	configText := found.ConfigText
+	found.mu.Unlock()
+	if m.db != nil {
+		_ = m.db.SavePluginState(name, enabled, configText)
+	}
+}
+
+func (m *Manager) SaveConfig(name, configText string) {
+	m.mu.RLock()
+	var found *Plugin
+	for _, p := range m.plugins {
+		if p.Name == name {
+			found = p
+			break
+		}
+	}
+	m.mu.RUnlock()
+	if found == nil {
+		return
+	}
+	found.mu.Lock()
+	found.ConfigText = configText
+	enabled := found.Enabled
+	hc, hasOnStart := found.hooks["on_start"]
+	found.mu.Unlock()
+	if m.db != nil {
+		_ = m.db.SavePluginState(name, enabled, configText)
+	}
+	if !hasOnStart {
+		return
+	}
+	// Re-run on_start so the plugin can re-parse the new config.
+	if hc.Sync {
+		found.mu.Lock()
+		if _, err := callHook(found, "on_start", lua.LString(configText)); err != nil {
+			log.Printf("plugin %s on_start (config reload): %v", name, err)
+		}
+		found.mu.Unlock()
+	} else {
+		go func() {
+			found.mu.Lock()
+			if _, err := callHook(found, "on_start", lua.LString(configText)); err != nil {
+				log.Printf("plugin %s on_start (config reload): %v", name, err)
+			}
+			found.mu.Unlock()
+		}()
+	}
+}
+
+func (m *Manager) RunOnStart() {
+	for _, p := range m.GetPlugins() {
+		if !p.Enabled {
+			continue
+		}
+		if _, ok := p.hooks["on_start"]; !ok {
+			continue
+		}
+		p.mu.Lock()
+		if _, err := callHook(p, "on_start", lua.LString(p.ConfigText)); err != nil {
+			log.Printf("plugin %s on_start: %v", p.Name, err)
+		}
+		p.mu.Unlock()
+	}
+}
+
+func (m *Manager) RunOnQuit() {
+	for _, p := range m.GetPlugins() {
+		if !p.Enabled {
+			continue
+		}
+		if _, ok := p.hooks["on_quit"]; !ok {
+			continue
+		}
+		p.mu.Lock()
+		if _, err := callHook(p, "on_quit"); err != nil {
+			log.Printf("plugin %s on_quit: %v", p.Name, err)
+		}
+		p.mu.Unlock()
+	}
+}
+
+func (m *Manager) RunSyncOnRequest(f *goproxy.Flow) intercept.Decision {
+	for _, p := range m.GetPlugins() {
+		if !p.Enabled {
+			continue
+		}
+		hc, ok := p.hooks["on_request"]
+		if !ok || !hc.Sync {
+			continue
+		}
+		p.mu.Lock()
+		result, err := callHook(p, "on_request", pushRequest(p.L, f))
+		p.mu.Unlock()
+		if err != nil {
+			log.Printf("plugin %s on_request: %v", p.Name, err)
+			continue
+		}
+		switch result {
+		case "drop":
+			return intercept.Drop
+		case "forward":
+			return intercept.Forward
+		}
+	}
+	return intercept.Intercept
+}
+
+func (m *Manager) RunAsyncOnRequest(f *goproxy.Flow) {
+	for _, p := range m.GetPlugins() {
+		if !p.Enabled {
+			continue
+		}
+		hc, ok := p.hooks["on_request"]
+		if !ok || hc.Sync {
+			continue
+		}
+		go func(p *Plugin) {
+			p.mu.Lock()
+			if _, err := callHook(p, "on_request", pushRequest(p.L, f)); err != nil {
+				log.Printf("plugin %s on_request: %v", p.Name, err)
+			}
+			p.mu.Unlock()
+		}(p)
+	}
+}
+
+func (m *Manager) RunSyncOnResponse(f *goproxy.Flow) intercept.Decision {
+	for _, p := range m.GetPlugins() {
+		if !p.Enabled {
+			continue
+		}
+		hc, ok := p.hooks["on_response"]
+		if !ok || !hc.Sync {
+			continue
+		}
+		p.mu.Lock()
+		result, err := callHook(p, "on_response", pushRequest(p.L, f), pushResponse(p.L, f))
+		p.mu.Unlock()
+		if err != nil {
+			log.Printf("plugin %s on_response: %v", p.Name, err)
+			continue
+		}
+		switch result {
+		case "drop":
+			return intercept.Drop
+		case "forward":
+			return intercept.Forward
+		}
+	}
+	return intercept.Intercept
+}
+
+func (m *Manager) RunAsyncOnResponse(f *goproxy.Flow) {
+	for _, p := range m.GetPlugins() {
+		if !p.Enabled {
+			continue
+		}
+		hc, ok := p.hooks["on_response"]
+		if !ok || hc.Sync {
+			continue
+		}
+		go func(p *Plugin) {
+			p.mu.Lock()
+			if _, err := callHook(p, "on_response", pushRequest(p.L, f), pushResponse(p.L, f)); err != nil {
+				log.Printf("plugin %s on_response: %v", p.Name, err)
+			}
+			p.mu.Unlock()
+		}(p)
+	}
+}
+
+func (m *Manager) RunOnHistoryEntry(e db.Entry) {
+	for _, p := range m.GetPlugins() {
+		if !p.Enabled {
+			continue
+		}
+		if _, ok := p.hooks["on_history_entry"]; !ok {
+			continue
+		}
+		go func(p *Plugin) {
+			p.mu.Lock()
+			if _, err := callHook(p, "on_history_entry", pushEntry(p.L, e)); err != nil {
+				log.Printf("plugin %s on_history_entry: %v", p.Name, err)
+			}
+			p.mu.Unlock()
+		}(p)
+	}
+}

internal/plugins/types.go

@@ -0,0 +1,66 @@
+package plugins
+
+import (
+	"sync"
+
+	lua "github.com/yuin/gopher-lua"
+)
+
+type HookConfig struct {
+	Sync bool
+}
+
+type Plugin struct {
+	Name       string
+	FilePath   string
+	Enabled    bool
+	ConfigText string
+
+	L     *lua.LState
+	mu    sync.Mutex
+	hooks map[string]HookConfig
+}
+
+func (p *Plugin) HookNames() []string {
+	out := make([]string, 0, len(p.hooks))
+	for name := range p.hooks {
+		out = append(out, name)
+	}
+	return out
+}
+
+func (p *Plugin) HookConfig(name string) (HookConfig, bool) {
+	hc, ok := p.hooks[name]
+	return hc, ok
+}
+
+type Info struct {
+	Name       string
+	FilePath   string
+	Enabled    bool
+	ConfigText string
+	Hooks      map[string]HookConfig
+}
+
+func (p *Plugin) Info() Info {
+	hooks := make(map[string]HookConfig, len(p.hooks))
+	for k, v := range p.hooks {
+		hooks[k] = v
+	}
+	return Info{
+		Name:       p.Name,
+		FilePath:   p.FilePath,
+		Enabled:    p.Enabled,
+		ConfigText: p.ConfigText,
+		Hooks:      hooks,
+	}
+}
+
+type PluginNotifMsg struct {
+	Title string
+	Body  string
+}
+
+type PluginQuitMsg struct {
+	Reason string
+}