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