1package policy
  2
  3import (
  4	"bytes"
  5	"fmt"
  6	"io"
  7	"log"
  8	"os"
  9	"path"
 10	"slices"
 11	"strings"
 12	"testing"
 13
 14	"postern/internal/db"
 15
 16	"github.com/emersion/go-message/mail"
 17	"go.starlark.net/starlark"
 18)
 19
 20func TestExampleModule(t *testing.T) {
 21	posternDB, err := db.OpenDB(&db.Config{
 22		DBPath: path.Join(t.TempDir(), "meta.db"),
 23	})
 24	if err != nil {
 25		t.Fatal(err)
 26	}
 27
 28	defer func(posternDB *db.DB) {
 29		err := posternDB.Close()
 30		if err != nil {
 31			log.Fatal(err)
 32		}
 33	}(posternDB)
 34
 35	policy := `
 36def init():
 37    return None
 38
 39def example(user, msg, source_folder, dest_folder):
 40    filenames = []
 41    pdf_data = b""
 42    bin_data = ""
 43    for a in msg.attachments:
 44        filenames.append(a.filename)
 45        if a.filename == "sample.pdf":
 46            pdf_data = a.data
 47        elif a.filename == "data.bin":
 48            bin_data = a.data
 49    return {
 50        "subject": msg.subject,
 51        "from": msg.header_from.address,
 52        "to": msg.header_to.address,
 53        "text_body": msg.text_body,
 54        "html_body": msg.html_body,
 55        "attachment_count": len(msg.attachments),
 56        "filenames": filenames,
 57        "pdf_data": pdf_data,
 58        "bin_data": bin_data,
 59    }
 60`
 61
 62	policyFile := path.Join(t.TempDir(), "policy.star")
 63	err = os.WriteFile(policyFile, []byte(policy), 0644)
 64	if err != nil {
 65		t.Fatal(err)
 66	}
 67
 68	e, err := NewEngine(policyFile, posternDB)
 69	if err != nil {
 70		t.Fatal(err)
 71	}
 72	s := e.state.Load()
 73
 74	hookVal, ok := s.globals["example"]
 75	if !ok {
 76		t.Fatal("example not found in engine globals")
 77	}
 78
 79	msgBody, err := testMessage.assemble()
 80	if err != nil {
 81		t.Fatal(err)
 82	}
 83
 84	m := Message{
 85		Size: 0,
 86		OpenStream: func() (io.ReadCloser, error) {
 87			return io.NopCloser(bytes.NewReader(msgBody)), nil
 88		},
 89	}
 90
 91	k, v := starlark.String("key"), starlark.MakeInt(1)
 92	result, err := starlark.Call(s.thread, hookVal, starlark.Tuple{k, &m, k, v}, nil)
 93	if err != nil {
 94		t.Fatal(err)
 95	}
 96
 97	// Assert the returned dict contains the expected values
 98	dict, ok := result.(*starlark.Dict)
 99	if !ok {
100		t.Fatalf("expected dict result, got %s", result.Type())
101	}
102
103	// Helper to get bytes value from dict
104	dictBytes := func(dict *starlark.Dict, key string) []byte {
105		v, _, err := dict.Get(starlark.String(key))
106		if err != nil || v == starlark.None {
107			return nil
108		}
109		if b, ok := v.(starlark.Bytes); ok {
110			return []byte(b)
111		}
112		return nil
113	}
114
115	// Helper to get list value from dict
116	dictList := func(dict *starlark.Dict, key string) []string {
117		v, _, err := dict.Get(starlark.String(key))
118		if err != nil || v == starlark.None {
119			return nil
120		}
121		if l, ok := v.(*starlark.List); ok {
122			r := make([]string, l.Len())
123			for i := 0; i < l.Len(); i++ {
124				item := l.Index(i)
125				if s, ok := item.(starlark.String); ok {
126					r[i] = string(s)
127				}
128			}
129			return r
130		}
131		return nil
132	}
133
134	// Validate each field
135	if got := dictString(dict, "subject"); got != testMessage.subject {
136		t.Errorf("subject = %q, want %q", got, testMessage.subject)
137	}
138
139	if got := dictString(dict, "from"); got != testMessage.fromAddress {
140		t.Errorf("from = %q, want %q", got, testMessage.fromAddress)
141	}
142
143	if got := dictString(dict, "to"); got != testMessage.toAddress {
144		t.Errorf("to = %q, want %q", got, testMessage.toAddress)
145	}
146
147	expectedTextBody := strings.ReplaceAll(testMessage.textBody, "\n", "\r\n")
148	if got := dictString(dict, "text_body"); got != expectedTextBody {
149		t.Errorf("text_body = %q, want %q", got, expectedTextBody)
150	}
151
152	if got := dictString(dict, "html_body"); got != testMessage.htmlBody {
153		t.Errorf("html_body = %q, want %q", got, testMessage.htmlBody)
154	}
155
156	if got := dictInt(dict, "attachment_count"); got != len(testMessage.attachments) {
157		t.Errorf("attachment_count = %d, want %d", got, len(testMessage.attachments))
158	}
159
160	filenames := dictList(dict, "filenames")
161	if len(filenames) != len(testMessage.attachments) {
162		t.Errorf("filenames length = %d, want %d", len(filenames), len(testMessage.attachments))
163	}
164
165	for _, f := range testMessage.attachments {
166		if !slices.Contains(filenames, f.filename) {
167			t.Errorf("filenames = %q, want %q", filenames, f.filename)
168		}
169	}
170
171	if got := dictBytes(dict, "pdf_data"); string(got) != string(testMessage.attachments[0].data) {
172		t.Errorf("pdf_data = %v, want %v", got, testMessage.attachments[0].data)
173	}
174
175	if got := dictBytes(dict, "bin_data"); string(got) != string(testMessage.attachments[1].data) {
176		t.Errorf("bin_data = %v, want %v", got, testMessage.attachments[1].data)
177	}
178}
179
180type attachmentFixture struct {
181	filename    string
182	contentType string
183	data        []byte
184}
185
186type messageFixture struct {
187	subject     string
188	fromHeader  string
189	fromAddress string
190	toAddress   string
191	textBody    string
192	htmlBody    string
193	attachments []attachmentFixture
194}
195
196var testMessage = messageFixture{
197	subject:     "Test Subject",
198	fromHeader:  `"noreply@example.com" <noreply@example.com>`,
199	fromAddress: "noreply@example.com",
200	toAddress:   "postern@example.com",
201	textBody:    "Hello,\n\nThis is the plain text body of the email.",
202	htmlBody:    "<html><body><h1>Hello</h1><p>This is the <b>HTML</b> body of the email.</p></body></html>",
203	attachments: []attachmentFixture{
204		{filename: "sample.pdf", contentType: "application/pdf", data: []byte("%PDF-1.4\n1 0 obj\n<< /Type /Catalog >>\nendobj\n%%EOF")},
205		{filename: "data.bin", contentType: "application/octet-stream", data: []byte{0x7F, 0x45, 0x4C, 0x46, 0x02, 0x01, 0x01, 0x00, 0x00}},
206	},
207}
208
209func (f messageFixture) assemble() ([]byte, error) {
210	headers := map[string][]string{
211		"Delivered-To": {"postern@example.com"},
212		"Received":     {"by mx1.p25.dev with ESMTP id b9d16c6997d311f18ac19200070e8571; Fri, 14 Aug 2026 11:31:44 +0000"},
213		"Date":         {"Fri, 14 Aug 2026 11:31:43 +0000 (UTC)"},
214		"From":         {f.fromHeader},
215		"To":           {f.toAddress},
216		"Message-ID":   {"<8555705.115511.1786707553174@example.com>"},
217		"Subject":      {f.subject},
218	}
219
220	var buf bytes.Buffer
221
222	// Create top-level writer (defaults to multipart/mixed when inline/attachment parts are added)
223	mw, err := mail.CreateWriter(&buf, mail.HeaderFromMap(headers))
224	if err != nil {
225		return nil, fmt.Errorf("failed to create mail writer: %w", err)
226	}
227
228	inlineWriter, err := mw.CreateInline()
229	if err != nil {
230		return nil, fmt.Errorf("failed to create inline writer: %w", err)
231	}
232
233	// Plain Text Part
234	var textHeader mail.InlineHeader
235	textHeader.SetContentType("text/plain", map[string]string{"charset": "utf-8"})
236	textWriter, err := inlineWriter.CreatePart(textHeader)
237	if err != nil {
238		return nil, fmt.Errorf("failed to create text part: %w", err)
239	}
240	if _, err := textWriter.Write([]byte(f.textBody)); err != nil {
241		return nil, err
242	}
243	if err := textWriter.Close(); err != nil {
244		return nil, err
245	}
246
247	// HTML Part
248	var htmlHeader mail.InlineHeader
249	htmlHeader.SetContentType("text/html", map[string]string{"charset": "utf-8"})
250	htmlWriter, err := inlineWriter.CreatePart(htmlHeader)
251	if err != nil {
252		return nil, fmt.Errorf("failed to create HTML part: %w", err)
253	}
254	if _, err := htmlWriter.Write([]byte(f.htmlBody)); err != nil {
255		return nil, err
256	}
257	if err := htmlWriter.Close(); err != nil {
258		return nil, err
259	}
260
261	// Close alternative container
262	if err := inlineWriter.Close(); err != nil {
263		return nil, err
264	}
265
266	// Attachments
267	for _, att := range f.attachments {
268		var header mail.AttachmentHeader
269		header.SetContentType(att.contentType, nil)
270		header.SetFilename(att.filename)
271		writer, err := mw.CreateAttachment(header)
272		if err != nil {
273			return nil, fmt.Errorf("failed to create attachment %s: %w", att.filename, err)
274		}
275		if _, err := writer.Write(att.data); err != nil {
276			return nil, err
277		}
278		if err := writer.Close(); err != nil {
279			return nil, err
280		}
281	}
282
283	// Finalize email body formatting and boundaries
284	if err := mw.Close(); err != nil {
285		return nil, err
286	}
287
288	return buf.Bytes(), nil
289}