email-mcp/internal/imapclient/backend.go

package imapclient

import (
	"context"
	"errors"
	"net"
	"strings"
	"time"

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

	"email-mcp/internal/secretstore"
)

const (
	defaultListMessagesLimit = 20
	maxListMessagesLimit     = 50
	imapImplicitTLSPort      = "993"
)

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

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
		}

		summaries = make([]MessageSummary, 0, len(fetched))
		for i := len(fetched) - 1; i >= 0; i-- {
			msg := fetched[i]
			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
			}
			message = Message{
				UID:     uid,
				Mailbox: mailbox,
				Headers: parseHeaderBlock(item.FindBodySection(headerSection)),
				Body:    string(item.FindBodySection(textSection)),
			}
			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 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 first.Name
	}
	if strings.TrimSpace(first.Name) == "" {
		return address
	}
	return first.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
}

var _ Backend = imapBackend{}