proton-bridge/internal/user/imap.go

package user

import (
	"context"
	"crypto/subtle"
	"fmt"
	"strings"
	"time"

	"github.com/ProtonMail/gluon/imap"
	"github.com/bradenaw/juniper/xslices"
	"gitlab.protontech.ch/go/liteapi"
	"golang.org/x/exp/slices"
)

var (
	defaultFlags          = imap.NewFlagSet(imap.FlagSeen, imap.FlagFlagged, imap.FlagDeleted)
	defaultPermanentFlags = imap.NewFlagSet(imap.FlagSeen, imap.FlagFlagged, imap.FlagDeleted)
	defaultAttributes     = imap.NewFlagSet()
)

const (
	folderPrefix = "Folders"
	labelPrefix  = "Labels"
)

type imapConnector struct {
	addrID   string
	client   *liteapi.Client
	updateCh <-chan imap.Update

	emails   []string
	password []byte

	flags, permFlags, attrs imap.FlagSet
}

func newIMAPConnector(
	client *liteapi.Client,
	updateCh <-chan imap.Update,
	password []byte,
	emails ...string,
) *imapConnector {
	return &imapConnector{
		client:   client,
		updateCh: updateCh,

		emails:   emails,
		password: password,

		flags:     defaultFlags,
		permFlags: defaultPermanentFlags,
		attrs:     defaultAttributes,
	}
}

// Authorize returns whether the given username/password combination are valid for this connector.
func (conn *imapConnector) Authorize(username string, password []byte) bool {
	if subtle.ConstantTimeCompare(conn.password, password) != 1 {
		return false
	}

	return xslices.IndexFunc(conn.emails, func(address string) bool {
		return strings.EqualFold(address, username)
	}) >= 0
}

// GetLabel returns information about the label with the given ID.
func (conn *imapConnector) GetLabel(ctx context.Context, labelID imap.LabelID) (imap.Mailbox, error) {
	label, err := conn.client.GetLabel(ctx, string(labelID), liteapi.LabelTypeLabel, liteapi.LabelTypeFolder)
	if err != nil {
		return imap.Mailbox{}, err
	}

	var name []string

	switch label.Type {
	case liteapi.LabelTypeLabel:
		name = []string{labelPrefix, label.Name}

	case liteapi.LabelTypeFolder:
		name = []string{folderPrefix, label.Name}

	default:
		name = []string{label.Name}
	}

	return imap.Mailbox{
		ID:             imap.LabelID(label.ID),
		Name:           name,
		Flags:          conn.flags,
		PermanentFlags: conn.permFlags,
		Attributes:     conn.attrs,
	}, nil
}

// CreateLabel creates a label with the given name.
func (conn *imapConnector) CreateLabel(ctx context.Context, name []string) (imap.Mailbox, error) {
	if len(name) != 2 {
		panic("subfolders are unsupported")
	}

	var labelType liteapi.LabelType

	if name[0] == folderPrefix {
		labelType = liteapi.LabelTypeFolder
	} else {
		labelType = liteapi.LabelTypeLabel
	}

	label, err := conn.client.CreateLabel(ctx, liteapi.CreateLabelReq{
		Name:  name[1:][0],
		Color: "#f66",
		Type:  labelType,
	})
	if err != nil {
		return imap.Mailbox{}, err
	}

	return imap.Mailbox{
		ID:             imap.LabelID(label.ID),
		Name:           name,
		Flags:          conn.flags,
		PermanentFlags: conn.permFlags,
		Attributes:     conn.attrs,
	}, nil
}

// UpdateLabel sets the name of the label with the given ID.
func (conn *imapConnector) UpdateLabel(ctx context.Context, labelID imap.LabelID, newName []string) error {
	if len(newName) != 2 {
		panic("subfolders are unsupported")
	}

	label, err := conn.client.GetLabel(ctx, string(labelID), liteapi.LabelTypeLabel, liteapi.LabelTypeFolder)
	if err != nil {
		return err
	}

	switch label.Type {
	case liteapi.LabelTypeFolder:
		if newName[0] != folderPrefix {
			return fmt.Errorf("cannot rename folder to label")
		}

	case liteapi.LabelTypeLabel:
		if newName[0] != labelPrefix {
			return fmt.Errorf("cannot rename label to folder")
		}

	case liteapi.LabelTypeSystem:
		return fmt.Errorf("cannot rename system label %q", label.Name)
	}

	if _, err := conn.client.UpdateLabel(ctx, label.ID, liteapi.UpdateLabelReq{
		Name:  newName[1:][0],
		Color: label.Color,
	}); err != nil {
		return err
	}

	return nil
}

// DeleteLabel deletes the label with the given ID.
func (conn *imapConnector) DeleteLabel(ctx context.Context, labelID imap.LabelID) error {
	return conn.client.DeleteLabel(ctx, string(labelID))
}

// GetMessage returns the message with the given ID.
func (conn *imapConnector) GetMessage(ctx context.Context, messageID imap.MessageID) (imap.Message, []imap.LabelID, error) {
	message, err := conn.client.GetMessage(ctx, string(messageID))
	if err != nil {
		return imap.Message{}, nil, err
	}

	flags := imap.NewFlagSet()

	if !message.Unread {
		flags = flags.Add(imap.FlagSeen)
	}

	if slices.Contains(message.LabelIDs, liteapi.StarredLabel) {
		flags = flags.Add(imap.FlagFlagged)
	}

	return imap.Message{
		ID:    imap.MessageID(message.ID),
		Flags: flags,
		Date:  time.Unix(message.Time, 0),
	}, mapTo[string, imap.LabelID](message.LabelIDs), nil
}

// CreateMessage creates a new message on the remote.
func (conn *imapConnector) CreateMessage(
	ctx context.Context,
	labelID imap.LabelID,
	literal []byte,
	parsedMessage *imap.ParsedMessage,
	flags imap.FlagSet,
	date time.Time,
) (imap.Message, error) {
	return imap.Message{}, ErrNotImplemented
}

// LabelMessages labels the given messages with the given label ID.
func (conn *imapConnector) LabelMessages(ctx context.Context, messageIDs []imap.MessageID, labelID imap.LabelID) error {
	return conn.client.LabelMessages(ctx, mapTo[imap.MessageID, string](messageIDs), string(labelID))
}

// UnlabelMessages unlabels the given messages with the given label ID.
func (conn *imapConnector) UnlabelMessages(ctx context.Context, messageIDs []imap.MessageID, labelID imap.LabelID) error {
	return conn.client.UnlabelMessages(ctx, mapTo[imap.MessageID, string](messageIDs), string(labelID))
}

// MoveMessages removes the given messages from one label and adds them to the other label.
func (conn *imapConnector) MoveMessages(ctx context.Context, messageIDs []imap.MessageID, labelFromID imap.LabelID, labelToID imap.LabelID) error {
	if err := conn.client.LabelMessages(ctx, mapTo[imap.MessageID, string](messageIDs), string(labelToID)); err != nil {
		return fmt.Errorf("labeling messages: %w", err)
	}

	if err := conn.client.UnlabelMessages(ctx, mapTo[imap.MessageID, string](messageIDs), string(labelFromID)); err != nil {
		return fmt.Errorf("unlabeling messages: %w", err)
	}

	return nil
}

// MarkMessagesSeen sets the seen value of the given messages.
func (conn *imapConnector) MarkMessagesSeen(ctx context.Context, messageIDs []imap.MessageID, seen bool) error {
	if seen {
		return conn.client.MarkMessagesRead(ctx, mapTo[imap.MessageID, string](messageIDs)...)
	} else {
		return conn.client.MarkMessagesUnread(ctx, mapTo[imap.MessageID, string](messageIDs)...)
	}
}

// MarkMessagesFlagged sets the flagged value of the given messages.
func (conn *imapConnector) MarkMessagesFlagged(ctx context.Context, messageIDs []imap.MessageID, flagged bool) error {
	if flagged {
		return conn.client.LabelMessages(ctx, mapTo[imap.MessageID, string](messageIDs), liteapi.StarredLabel)
	} else {
		return conn.client.UnlabelMessages(ctx, mapTo[imap.MessageID, string](messageIDs), liteapi.StarredLabel)
	}
}

// GetUpdates returns a stream of updates that the gluon server should apply.
// It is recommended that the returned channel is buffered with at least constants.ChannelBufferCount.
func (conn *imapConnector) GetUpdates() <-chan imap.Update {
	return conn.updateCh
}

// GetUIDValidity returns the default UID validity for this user.
func (conn *imapConnector) GetUIDValidity() imap.UID {
	return imap.UID(1)
}

// SetUIDValidity sets the default UID validity for this user.
func (conn *imapConnector) SetUIDValidity(uidValidity imap.UID) error {
	return nil
}

// Close the connector will no longer be used and all resources should be closed/released.
func (conn *imapConnector) Close(ctx context.Context) error {
	return nil
}