email-mcp/internal/imapclient/backend.go

package imapclient

import (
	"bytes"
	"context"
	"errors"
	"io"
	"mime"
	"net"
	"sort"
	"strings"
	"time"

	goimap "github.com/emersion/go-imap/v2"
	goimapclient "github.com/emersion/go-imap/v2/imapclient"
	gomail "github.com/emersion/go-message/mail"

	"email-mcp/internal/secretstore"
)

const (
	DefaultListMessagesLimit = 20
	MaxListMessagesLimit     = 50
	imapImplicitTLSPort      = "993"
)

var ErrMessageNotFound = errors.New("message not found")

var headerWordDecoder = mime.WordDecoder{}

type waitCommand interface {
	Wait() error
}

type listCommand interface {
	Collect() ([]*goimap.ListData, error)
}

type selectCommand interface {
	Wait() (*goimap.SelectData, error)
}

type fetchCommand interface {
	Collect() ([]*goimapclient.FetchMessageBuffer, error)
}

type imapSession interface {
	Login(username, password string) waitCommand
	Logout() waitCommand
	Close() error
	List(ref, pattern string, options *goimap.ListOptions) listCommand
	Select(mailbox string, options *goimap.SelectOptions) selectCommand
	Fetch(numSet goimap.NumSet, options *goimap.FetchOptions) fetchCommand
}

type imapBackend struct {
	dialTLS func(address string, options *goimapclient.Options) (imapSession, error)
}

type goIMAPSession struct {
	client *goimapclient.Client
}

func NewDefaultBackend() Backend {
	return imapBackend{
		dialTLS: func(address string, options *goimapclient.Options) (imapSession, error) {
			client, err := goimapclient.DialTLS(address, options)
			if err != nil {
				return nil, err
			}
			return goIMAPSession{client: client}, nil
		},
	}
}

func (s goIMAPSession) Login(username, password string) waitCommand {
	return s.client.Login(username, password)
}

func (s goIMAPSession) Logout() waitCommand {
	return s.client.Logout()
}

func (s goIMAPSession) Close() error {
	return s.client.Close()
}

func (s goIMAPSession) List(ref, pattern string, options *goimap.ListOptions) listCommand {
	return s.client.List(ref, pattern, options)
}

func (s goIMAPSession) Select(mailbox string, options *goimap.SelectOptions) selectCommand {
	return s.client.Select(mailbox, options)
}

func (s goIMAPSession) Fetch(numSet goimap.NumSet, options *goimap.FetchOptions) fetchCommand {
	return s.client.Fetch(numSet, options)
}

func (b imapBackend) ListMailboxes(ctx context.Context, cred secretstore.Credential) ([]Mailbox, error) {
	var mailboxes []Mailbox
	err := b.withSession(ctx, cred, func(session imapSession) error {
		listed, err := session.List("", "*", nil).Collect()
		if err != nil {
			return err
		}

		mailboxes = make([]Mailbox, 0, len(listed))
		for _, item := range listed {
			if item == nil || hasMailboxAttr(item.Attrs, goimap.MailboxAttrNoSelect) {
				continue
			}
			mailboxes = append(mailboxes, Mailbox{Name: item.Mailbox})
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	return mailboxes, nil
}

func (b imapBackend) ListMessages(ctx context.Context, cred secretstore.Credential, mailbox string, limit int) ([]MessageSummary, error) {
	var summaries []MessageSummary
	err := b.withSession(ctx, cred, func(session imapSession) error {
		selected, err := session.Select(mailbox, &goimap.SelectOptions{ReadOnly: true}).Wait()
		if err != nil {
			return err
		}
		if selected == nil || selected.NumMessages == 0 {
			summaries = []MessageSummary{}
			return nil
		}

		effectiveLimit := clampListLimit(limit)
		start := uint32(1)
		if selected.NumMessages > uint32(effectiveLimit) {
			start = selected.NumMessages - uint32(effectiveLimit) + 1
		}

		var seqSet goimap.SeqSet
		seqSet.AddRange(start, selected.NumMessages)

		fetched, err := session.Fetch(seqSet, &goimap.FetchOptions{
			UID:      true,
			Envelope: true,
		}).Collect()
		if err != nil {
			return err
		}

		sort.SliceStable(fetched, func(i, j int) bool {
			if fetched[i] == nil {
				return false
			}
			if fetched[j] == nil {
				return true
			}
			if fetched[i].UID == fetched[j].UID {
				return fetched[i].SeqNum > fetched[j].SeqNum
			}
			return fetched[i].UID > fetched[j].UID
		})

		summaries = make([]MessageSummary, 0, len(fetched))
		for _, msg := range fetched {
			if msg == nil {
				continue
			}
			summaries = append(summaries, MessageSummary{
				UID:     uint32(msg.UID),
				Subject: envelopeSubject(msg.Envelope),
				From:    envelopeFrom(msg.Envelope),
			})
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	return summaries, nil
}

func (b imapBackend) GetMessage(ctx context.Context, cred secretstore.Credential, mailbox string, uid uint32) (Message, error) {
	var message Message
	err := b.withSession(ctx, cred, func(session imapSession) error {
		if _, err := session.Select(mailbox, &goimap.SelectOptions{ReadOnly: true}).Wait(); err != nil {
			return err
		}

		headerSection := &goimap.FetchItemBodySection{
			Specifier: goimap.PartSpecifierHeader,
			Peek:      true,
		}
		textSection := &goimap.FetchItemBodySection{
			Specifier: goimap.PartSpecifierText,
			Peek:      true,
		}

		fetched, err := session.Fetch(goimap.UIDSetNum(goimap.UID(uid)), &goimap.FetchOptions{
			UID:         true,
			BodySection: []*goimap.FetchItemBodySection{headerSection, textSection},
		}).Collect()
		if err != nil {
			return err
		}

		for _, item := range fetched {
			if item == nil || uint32(item.UID) != uid {
				continue
			}
			headerBytes := item.FindBodySection(headerSection)
			bodyBytes := item.FindBodySection(textSection)
			message = Message{
				UID:     uid,
				Mailbox: mailbox,
				Headers: parseHeaderBlock(headerBytes),
				Body:    decodeBodyText(headerBytes, bodyBytes),
			}
			return nil
		}
		return ErrMessageNotFound
	})
	if err != nil {
		return Message{}, err
	}
	return message, nil
}

func (b imapBackend) withSession(ctx context.Context, cred secretstore.Credential, fn func(imapSession) error) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	if err := cred.Validate(); err != nil {
		return err
	}

	dialTLS := b.dialTLS
	if dialTLS == nil {
		dialTLS = func(address string, options *goimapclient.Options) (imapSession, error) {
			client, err := goimapclient.DialTLS(address, options)
			if err != nil {
				return nil, err
			}
			return goIMAPSession{client: client}, nil
		}
	}

	session, err := dialTLS(imapAddress(cred.Host), &goimapclient.Options{
		Dialer: &net.Dialer{Timeout: 30 * time.Second},
	})
	if err != nil {
		return err
	}
	defer session.Close()

	if err := session.Login(cred.Username, cred.Password).Wait(); err != nil {
		return err
	}

	runErr := fn(session)
	logoutErr := session.Logout().Wait()
	if runErr != nil {
		return runErr
	}
	return logoutErr
}

func imapAddress(host string) string {
	host = strings.TrimSpace(host)
	if _, _, err := net.SplitHostPort(host); err == nil {
		return host
	}
	return net.JoinHostPort(host, imapImplicitTLSPort)
}

func clampListLimit(limit int) int {
	if limit <= 0 {
		return DefaultListMessagesLimit
	}
	if limit > MaxListMessagesLimit {
		return MaxListMessagesLimit
	}
	return limit
}

func hasMailboxAttr(attrs []goimap.MailboxAttr, want goimap.MailboxAttr) bool {
	for _, attr := range attrs {
		if attr == want {
			return true
		}
	}
	return false
}

func envelopeSubject(envelope *goimap.Envelope) string {
	if envelope == nil {
		return ""
	}
	return decodeEncodedWords(envelope.Subject)
}

func envelopeFrom(envelope *goimap.Envelope) string {
	if envelope == nil || len(envelope.From) == 0 {
		return ""
	}

	first := envelope.From[0]
	address := first.Addr()
	if address == "" {
		return decodeEncodedWords(first.Name)
	}
	name := decodeEncodedWords(first.Name)
	if strings.TrimSpace(name) == "" {
		return address
	}
	return name + " <" + address + ">"
}

func parseHeaderBlock(raw []byte) []Header {
	if len(raw) == 0 {
		return nil
	}

	lines := strings.Split(strings.ReplaceAll(string(raw), "\r\n", "\n"), "\n")
	headers := make([]Header, 0, len(lines))

	for _, line := range lines {
		if line == "" {
			break
		}

		if line[0] == ' ' || line[0] == '\t' {
			if len(headers) == 0 {
				continue
			}
			headers[len(headers)-1].Value += " " + strings.TrimSpace(line)
			continue
		}

		name, value, ok := strings.Cut(line, ":")
		if !ok {
			continue
		}
		headers = append(headers, Header{
			Name:  strings.TrimSpace(name),
			Value: strings.TrimSpace(value),
		})
	}

	return headers
}

func decodeEncodedWords(value string) string {
	decoded, err := headerWordDecoder.DecodeHeader(value)
	if err != nil {
		return value
	}
	return decoded
}

func decodeBodyText(headerBytes, bodyBytes []byte) string {
	if len(bodyBytes) == 0 {
		return ""
	}

	messageBytes := assembleMessage(headerBytes, bodyBytes)
	reader, err := gomail.CreateReader(bytes.NewReader(messageBytes))
	if err != nil && reader == nil {
		return string(bodyBytes)
	}
	if reader == nil {
		return string(bodyBytes)
	}

	var fallback string
	for {
		part, partErr := reader.NextPart()
		if errors.Is(partErr, io.EOF) {
			break
		}
		if partErr != nil && part == nil {
			break
		}
		if part == nil {
			continue
		}

		inlineHeader, ok := part.Header.(*gomail.InlineHeader)
		if !ok {
			continue
		}

		partBytes, readErr := io.ReadAll(part.Body)
		if readErr != nil {
			continue
		}

		mediaType, _, _ := inlineHeader.ContentType()
		switch strings.ToLower(mediaType) {
		case "text/plain":
			return string(partBytes)
		case "text/html":
			if fallback == "" {
				fallback = string(partBytes)
			}
		default:
			if fallback == "" && strings.HasPrefix(strings.ToLower(mediaType), "text/") {
				fallback = string(partBytes)
			}
		}
	}

	if fallback != "" {
		return fallback
	}
	return string(bodyBytes)
}

func assembleMessage(headerBytes, bodyBytes []byte) []byte {
	trimmedHeaders := bytes.TrimRight(headerBytes, "\r\n")
	assembled := make([]byte, 0, len(trimmedHeaders)+len(bodyBytes)+4)
	assembled = append(assembled, trimmedHeaders...)
	assembled = append(assembled, '\r', '\n', '\r', '\n')
	assembled = append(assembled, bodyBytes...)
	return assembled
}

var _ Backend = imapBackend{}