1package policy
  2
  3import (
  4	"context"
  5	"fmt"
  6	"log"
  7	"net/netip"
  8	"os"
  9
 10	"postern/internal/db"
 11
 12	"go.starlark.net/lib/json"
 13	sltime "go.starlark.net/lib/time"
 14	"go.starlark.net/starlark"
 15	"go.starlark.net/starlarkstruct"
 16	"go.starlark.net/syntax"
 17)
 18
 19type RelayConfig struct {
 20	Host     string
 21	Port     int
 22	Username string
 23	Password string
 24}
 25
 26type Engine struct {
 27	Listeners   []ListenerConfig
 28	UsersModule *UsersModule
 29	globals     starlark.StringDict
 30	thread      *starlark.Thread
 31}
 32
 33type ListenerConfig struct {
 34	Address     netip.AddrPort
 35	Role        string
 36	TLSCertFile string
 37	TLSKeyFile  string
 38}
 39
 40func NewEngine(policyFile string, database *db.DB) (Engine, error) {
 41	script, err := os.ReadFile(policyFile)
 42	if err != nil {
 43		log.Fatal(err)
 44	}
 45
 46	// Create an 'os' module containing our getenv function
 47	osModule := &starlarkstruct.Module{
 48		Name: "os",
 49		Members: starlark.StringDict{
 50			"getenv": starlark.NewBuiltin("getenv", starlarkGetenv),
 51		},
 52	}
 53
 54	// Create the users module
 55	usersModule := NewUsersModule(database)
 56
 57	// Expose the "os", "mda", "users", and "http" modules to the Starlark global environment
 58	m := &MDA{
 59		parsedListeners: make([]ListenerConfig, 0),
 60	}
 61
 62	predeclared := starlark.StringDict{
 63		"exec":  &ExecModule{},
 64		"http":  &HTTPModule{},
 65		"json":  json.Module,
 66		"kv":    NewKVModule(database),
 67		"mda":   m,
 68		"os":    osModule,
 69		"time":  sltime.Module,
 70		"users": usersModule,
 71	}
 72
 73	// Create a Starlark thread. We customize the Print function
 74	// so the Starlark `print()` built-in logs cleanly in Go.
 75	thread := &starlark.Thread{
 76		Name: "mda-main",
 77		Print: func(_ *starlark.Thread, msg string) {
 78			fmt.Printf("[init.star]: %s\n", msg)
 79		},
 80	}
 81	thread.SetLocal("context", context.Background())
 82
 83	// 1. Evaluate the script
 84	globals, err := starlark.ExecFileOptions(&syntax.FileOptions{}, thread, "policy.star", script, predeclared)
 85	if err != nil {
 86		return Engine{}, err
 87	}
 88
 89	// 2. Call the init() function
 90	initVal, ok := globals["init"]
 91	if !ok {
 92		return Engine{}, fmt.Errorf("init not found in policy")
 93	}
 94
 95	_, err = starlark.Call(thread, initVal, nil, nil)
 96	if err != nil {
 97		return Engine{}, err
 98	}
 99
100	engine := Engine{
101		Listeners:   m.parsedListeners,
102		UsersModule: usersModule,
103		globals:     globals,
104		thread:      thread,
105	}
106	return engine, err
107}
108
109func (e *Engine) OnMessageSubmit(user User, msg *Message) (RelayConfig, error) {
110	hookVal, ok := e.globals["on_message_submit"]
111	if !ok {
112		return RelayConfig{}, nil
113	}
114
115	result, err := starlark.Call(e.thread, hookVal, starlark.Tuple{user, msg}, nil)
116	if err != nil {
117		return RelayConfig{}, err
118	}
119
120	if result == starlark.None {
121		return RelayConfig{}, nil
122	}
123
124	dict, ok := result.(*starlark.Dict)
125	if !ok {
126		return RelayConfig{}, fmt.Errorf("on_message_submit must return a dict, got %s", result.Type())
127	}
128
129	relay := RelayConfig{
130		Host:     dictString(dict, "host"),
131		Port:     dictInt(dict, "port"),
132		Username: dictString(dict, "username"),
133		Password: dictString(dict, "password"),
134	}
135	return relay, nil
136}
137
138func (e *Engine) OnMessageDeliver(conn Conn, user User, msg *Message) (string, error) {
139	defaultMailbox := "INBOX"
140	hookVal, ok := e.globals["on_message_deliver"]
141	if !ok {
142		return defaultMailbox, nil
143	}
144
145	result, err := starlark.Call(e.thread, hookVal, starlark.Tuple{conn, user, msg}, nil)
146	if err != nil {
147		return defaultMailbox, err
148	}
149
150	mailbox := defaultMailbox
151	if result != starlark.None {
152		if s, ok := result.(starlark.String); ok {
153			mailbox = string(s)
154		}
155	}
156
157	return mailbox, nil
158}
159
160func (e *Engine) OnMessageMoved(user User, msg *Message, sourceFolder, destFolder string) error {
161	hookVal, ok := e.globals["on_message_moved"]
162	if !ok {
163		return nil
164	}
165
166	_, err := starlark.Call(e.thread, hookVal, starlark.Tuple{user, msg, starlark.String(sourceFolder), starlark.String(destFolder)}, nil)
167	return err
168}
169
170// starlarkGetenv implements Python's os.getenv("KEY", "default")
171func starlarkGetenv(_ *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
172	var key string
173	var def starlark.Value = starlark.None
174
175	if err := starlark.UnpackArgs(b.Name(), args, kwargs, "key", &key, "default?", &def); err != nil {
176		return nil, err
177	}
178
179	if val, exists := os.LookupEnv(key); exists {
180		return starlark.String(val), nil
181	}
182	return def, nil
183}
184
185func dictString(dict *starlark.Dict, key string) string {
186	v, _, err := dict.Get(starlark.String(key))
187	if err != nil || v == starlark.None {
188		return ""
189	}
190	if s, ok := v.(starlark.String); ok {
191		return string(s)
192	}
193	return ""
194}
195
196func dictInt(dict *starlark.Dict, key string) int {
197	v, _, err := dict.Get(starlark.String(key))
198	if err != nil || v == starlark.None {
199		return 0
200	}
201	if i, ok := v.(starlark.Int); ok {
202		if n, ok := i.Int64(); ok {
203			return int(n)
204		}
205	}
206	return 0
207}