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