kittemail/message.go

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"github.com/google/uuid"
	"io"
	"mime"
	"strings"
	"time"

	"github.com/emersion/go-imap/v2"
	"github.com/emersion/go-imap/v2/imapserver"
	gomessage "github.com/emersion/go-message"
	"github.com/emersion/go-message/mail"
	"github.com/emersion/go-message/textproto"
)

type message struct {
	// immutable
	uid   imap.UID
	buf   io.ReadCloser
	len   int64
	t     time.Time
	id    uuid.UUID
	lines int64

	// mutable, protected by Mailbox.mutex
	flags map[imap.Flag]struct{}
}

func (msg *message) fetch(w *imapserver.FetchResponseWriter, options *imap.FetchOptions) error {
	w.WriteUID(msg.uid)

	if options.Flags {
		fmt.Println("FLAGS: ", msg.flagList())
		w.WriteFlags(msg.flagList())
	}
	if options.InternalDate {
		w.WriteInternalDate(msg.t)
	}
	if options.RFC822Size {
		fmt.Println("RFC822SIZE: ", msg.len)
		w.WriteRFC822Size(msg.len)
	}
	if options.Envelope {
		w.WriteEnvelope(msg.envelope())
	}
	bs := options.BodyStructure
	if bs != nil {
		structure, err := msg.bodyStructure(bs.Extended)
		if err != nil {
			println("BODYSTRUCTURE ERROR: ", err.Error())
			return err
		}
		println("BODYSTRUCTURE: ", structure.Disposition().Value)
		w.WriteBodyStructure(structure)
	}

	for _, bs := range options.BodySection {
		buf := msg.bodySection(bs)
		wc := w.WriteBodySection(bs, msg.len)
		_, writeErr := wc.Write(buf)
		closeErr := wc.Close()
		if writeErr != nil {
			return writeErr
		}
		if closeErr != nil {
			return closeErr
		}
	}

	return w.Close()
}

func (msg *message) envelope() *imap.Envelope {
	br := bufio.NewReader(msg.buf)
	header, err := textproto.ReadHeader(br)
	if err != nil {
		return nil
	}
	return getEnvelope(header)
}

func (msg *message) bodyStructure(extended bool) (imap.BodyStructure, error) {
	br := bufio.NewReader(msg.buf)
	header, err := textproto.ReadHeader(br)
	if err != nil {
		return nil, err
	}
	return getBodyStructure(header, br, extended, uint32(msg.len), msg.lines)
}

func openMessagePart(header textproto.Header, body io.Reader, parentMediaType string) (textproto.Header, io.Reader) {
	msgHeader := gomessage.Header{Header: header}
	mediaType, _, _ := msgHeader.ContentType()
	if !msgHeader.Has("Content-Type") && parentMediaType == "multipart/digest" {
		mediaType = "message/rfc822"
	}
	if mediaType == "message/rfc822" || mediaType == "message/global" {
		br := bufio.NewReader(body)
		header, _ = textproto.ReadHeader(br)
		return header, br
	}
	return header, body
}

func (msg *message) bodySection(item *imap.FetchItemBodySection) []byte {
	var (
		header textproto.Header
		body   io.Reader
	)

	br := bufio.NewReader(msg.buf)
	header, err := textproto.ReadHeader(br)
	if err != nil {
		return nil
	}
	body = br

	// First part of non-multipart message refers to the message itself
	msgHeader := gomessage.Header{Header: header}
	mediaType, _, _ := msgHeader.ContentType()
	partPath := item.Part
	if !strings.HasPrefix(mediaType, "multipart/") && len(partPath) > 0 && partPath[0] == 1 {
		partPath = partPath[1:]
	}

	// Find the requested part using the provided path
	var parentMediaType string
	for i := 0; i < len(partPath); i++ {
		partNum := partPath[i]

		header, body = openMessagePart(header, body, parentMediaType)

		msgHeader := gomessage.Header{Header: header}
		mediaType, typeParams, _ := msgHeader.ContentType()
		if !strings.HasPrefix(mediaType, "multipart/") {
			if partNum != 1 {
				return nil
			}
			continue
		}

		mr := textproto.NewMultipartReader(body, typeParams["boundary"])
		found := false
		for j := 1; j <= partNum; j++ {
			p, err := mr.NextPart()
			if err != nil {
				return nil
			}

			if j == partNum {
				parentMediaType = mediaType
				header = p.Header
				body = p
				found = true
				break
			}
		}
		if !found {
			return nil
		}
	}

	if len(item.Part) > 0 {
		switch item.Specifier {
		case imap.PartSpecifierHeader, imap.PartSpecifierText:
			header, body = openMessagePart(header, body, parentMediaType)
		}
	}

	// Filter header fields
	if len(item.HeaderFields) > 0 {
		keep := make(map[string]struct{})
		for _, k := range item.HeaderFields {
			keep[strings.ToLower(k)] = struct{}{}
		}
		for field := header.Fields(); field.Next(); {
			if _, ok := keep[strings.ToLower(field.Key())]; !ok {
				field.Del()
			}
		}
	}
	for _, k := range item.HeaderFieldsNot {
		header.Del(k)
	}

	// Write the requested data to a buffer
	var buf bytes.Buffer

	writeHeader := true
	switch item.Specifier {
	case imap.PartSpecifierNone:
		writeHeader = len(item.Part) == 0
	case imap.PartSpecifierText:
		writeHeader = false
	}
	if writeHeader {
		if err := textproto.WriteHeader(&buf, header); err != nil {
			return nil
		}
	}

	switch item.Specifier {
	case imap.PartSpecifierNone, imap.PartSpecifierText:
		if _, err := io.Copy(&buf, body); err != nil {
			return nil
		}
	}

	// Extract partial if any
	b := buf.Bytes()
	if partial := item.Partial; partial != nil {
		end := partial.Offset + partial.Size
		if partial.Offset > int64(len(b)) {
			return nil
		}
		if end > int64(len(b)) {
			end = int64(len(b))
		}
		b = b[partial.Offset:end]
	}
	return b
}

func (msg *message) flagList() []imap.Flag {
	var flags []imap.Flag
	for flag := range msg.flags {
		flags = append(flags, flag)
	}
	return flags
}

func (msg *message) store(store *imap.StoreFlags) {
	switch store.Op {
	case imap.StoreFlagsSet:
		msg.flags = make(map[imap.Flag]struct{})
		fallthrough
	case imap.StoreFlagsAdd:
		for _, flag := range store.Flags {
			msg.flags[canonicalFlag(flag)] = struct{}{}
		}
	case imap.StoreFlagsDel:
		for _, flag := range store.Flags {
			delete(msg.flags, canonicalFlag(flag))
		}
	default:
		panic(fmt.Errorf("unknown STORE flag operation: %v", store.Op))
	}
}

func (msg *message) reader() *gomessage.Entity {
	r, _ := gomessage.Read(msg.buf)
	if r == nil {
		r, _ = gomessage.New(gomessage.Header{}, bytes.NewReader(nil))
	}
	return r
}

func (msg *message) search(seqNum uint32, criteria *imap.SearchCriteria) bool {
	for _, seqSet := range criteria.SeqNum {
		if seqNum == 0 || !seqSet.Contains(seqNum) {
			println("DOES NOT CONTAIN 1")
			return false
		}
	}
	for _, uidSet := range criteria.UID {
		if !uidSet.Contains(msg.uid) {
			println("DOES NOT CONTAIN 2")
			return false
		}
	}
	if !matchDate(msg.t, criteria.Since, criteria.Before) {
		println("DOES NOT CONTAIN 3")
		return false
	}

	for _, flag := range criteria.Flag {
		if _, ok := msg.flags[canonicalFlag(flag)]; !ok {
			println("DOES NOT CONTAIN 4")
			return false
		}
	}
	for _, flag := range criteria.NotFlag {
		if _, ok := msg.flags[canonicalFlag(flag)]; ok {
			println("DOES NOT CONTAIN 5")
			return false
		}
	}

	if criteria.Larger != 0 && msg.len <= criteria.Larger {
		println("DOES NOT CONTAIN 6")
		return false
	}
	if criteria.Smaller != 0 && msg.len >= criteria.Smaller {
		println("DOES NOT CONTAIN 7")
		return false
	}

	header := mail.Header{Header: msg.reader().Header}

	for _, fieldCriteria := range criteria.Header {
		if !matchHeaderFields(header.FieldsByKey(fieldCriteria.Key), fieldCriteria.Value) {
			println("DOES NOT CONTAIN 8")
			return false
		}
	}

	if !criteria.SentSince.IsZero() || !criteria.SentBefore.IsZero() {
		t, err := header.Date()
		if err != nil {
			println("DOES NOT CONTAIN 9")
			return false
		} else if !matchDate(t, criteria.SentSince, criteria.SentBefore) {
			println("DOES NOT CONTAIN 10")
			return false
		}
	}

	for _, text := range criteria.Text {
		if !matchEntity(msg.reader(), text, true) {
			println("DOES NOT CONTAIN 11")
			return false
		}
	}
	for _, body := range criteria.Body {
		if !matchEntity(msg.reader(), body, false) {
			println("DOES NOT CONTAIN 12")
			return false
		}
	}

	for _, not := range criteria.Not {
		if msg.search(seqNum, &not) {
			println("DOES NOT CONTAIN 13")
			return false
		}
	}
	for _, or := range criteria.Or {
		if !msg.search(seqNum, &or[0]) && !msg.search(seqNum, &or[1]) {
			println("DOES NOT CONTAIN 14")
			return false
		}
	}

	return true
}

func matchDate(t, since, before time.Time) bool {
	// We discard time zone information by setting it to UTC.
	// RFC 3501 explicitly requires zone unaware date comparison.
	t = time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, time.UTC)

	if !since.IsZero() && t.Before(since) {
		return false
	}
	if !before.IsZero() && !t.Before(before) {
		return false
	}
	return true
}

func matchHeaderFields(fields gomessage.HeaderFields, pattern string) bool {
	if pattern == "" {
		return fields.Len() > 0
	}

	pattern = strings.ToLower(pattern)
	for fields.Next() {
		v, _ := fields.Text()
		if strings.Contains(strings.ToLower(v), pattern) {
			return true
		}
	}
	return false
}

func matchEntity(e *gomessage.Entity, pattern string, includeHeader bool) bool {
	if pattern == "" {
		return true
	}

	if includeHeader && matchHeaderFields(e.Header.Fields(), pattern) {
		return true
	}

	if mr := e.MultipartReader(); mr != nil {
		for {
			part, err := mr.NextPart()
			if err == io.EOF {
				break
			} else if err != nil {
				return false
			}

			if matchEntity(part, pattern, includeHeader) {
				return true
			}
		}

		return false
	} else {
		t, _, err := e.Header.ContentType()
		if err != nil {
			return false
		}

		if !strings.HasPrefix(t, "text/") && !strings.HasPrefix(t, "message/") {
			return false
		}

		buf, err := io.ReadAll(e.Body)
		if err != nil {
			return false
		}

		return bytes.Contains(bytes.ToLower(buf), bytes.ToLower([]byte(pattern)))
	}
}

func getEnvelope(h textproto.Header) *imap.Envelope {
	mh := mail.Header{Header: gomessage.Header{Header: h}}
	date, _ := mh.Date()
	inReplyTo, _ := mh.MsgIDList("In-Reply-To")
	messageID, _ := mh.MessageID()
	return &imap.Envelope{
		Date:      date,
		Subject:   h.Get("Subject"),
		From:      parseAddressList(mh, "From"),
		Sender:    parseAddressList(mh, "Sender"),
		ReplyTo:   parseAddressList(mh, "Reply-To"),
		To:        parseAddressList(mh, "To"),
		Cc:        parseAddressList(mh, "Cc"),
		Bcc:       parseAddressList(mh, "Bcc"),
		InReplyTo: inReplyTo,
		MessageID: messageID,
	}
}

func parseAddressList(mh mail.Header, k string) []imap.Address {
	// TODO: leave the quoted words unchanged
	// TODO: handle groups
	addrs, _ := mh.AddressList(k)
	var l []imap.Address
	for _, addr := range addrs {
		mailbox, host, ok := strings.Cut(addr.Address, "@")
		if !ok {
			continue
		}
		l = append(l, imap.Address{
			Name:    mime.QEncoding.Encode("utf-8", addr.Name),
			Mailbox: mailbox,
			Host:    host,
		})
	}
	return l
}

func getBodyStructure(rawHeader textproto.Header, r io.Reader, extended bool, length uint32, lines int64) (imap.BodyStructure, error) {
	header := gomessage.Header{Header: rawHeader}

	mediaType, typeParams, _ := header.ContentType()
	primaryType, subType, _ := strings.Cut(mediaType, "/")

	if primaryType == "multipart" {
		bs := &imap.BodyStructureMultiPart{Subtype: subType}
		mr := textproto.NewMultipartReader(r, typeParams["boundary"])
		for {
			part, _ := mr.NextPart()
			if part == nil {
				break
			}
			structure, err := getBodyStructure(part.Header, part, extended, length, lines)
			if err != nil {
				return nil, err
			}
			bs.Children = append(bs.Children, structure)
		}
		if extended {
			bs.Extended = &imap.BodyStructureMultiPartExt{
				Params:      typeParams,
				Disposition: getContentDisposition(header),
				Language:    getContentLanguage(header),
				Location:    header.Get("Content-Location"),
			}
		}
		return bs, nil
	} else {
		bs := &imap.BodyStructureSinglePart{
			Type:        primaryType,
			Subtype:     subType,
			Params:      typeParams,
			ID:          header.Get("Content-Id"),
			Description: header.Get("Content-Description"),
			Encoding:    header.Get("Content-Transfer-Encoding"),
			Size:        length,
		}
		if mediaType == "message/rfc822" || mediaType == "message/global" {
			br := bufio.NewReader(r)
			childHeader, err := textproto.ReadHeader(br)
			if err != nil {
				return nil, err
			}
			structure, err := getBodyStructure(childHeader, br, extended, length, lines)
			if err != nil {
				return nil, err
			}
			bs.MessageRFC822 = &imap.BodyStructureMessageRFC822{
				Envelope:      getEnvelope(childHeader),
				BodyStructure: structure,
				NumLines:      lines,
			}
		}
		if primaryType == "text" {
			bs.Text = &imap.BodyStructureText{
				NumLines: lines,
			}
		}
		if extended {
			bs.Extended = &imap.BodyStructureSinglePartExt{
				Disposition: getContentDisposition(header),
				Language:    getContentLanguage(header),
				Location:    header.Get("Content-Location"),
			}
		}
		return bs, nil
	}
}

func getContentDisposition(header gomessage.Header) *imap.BodyStructureDisposition {
	disp, dispParams, _ := header.ContentDisposition()
	if disp == "" {
		return nil
	}
	return &imap.BodyStructureDisposition{
		Value:  disp,
		Params: dispParams,
	}
}

func getContentLanguage(header gomessage.Header) []string {
	v := header.Get("Content-Language")
	if v == "" {
		return nil
	}
	// TODO: handle CFWS
	l := strings.Split(v, ",")
	for i, lang := range l {
		l[i] = strings.TrimSpace(lang)
	}
	return l
}

func canonicalFlag(flag imap.Flag) imap.Flag {
	return imap.Flag(strings.ToLower(string(flag)))
}