Import/Export backend prep

internal/transfer/provider_pmapi.go

@@ -0,0 +1,140 @@
+// Copyright (c) 2020 Proton Technologies AG
+//
+// This file is part of ProtonMail Bridge.
+//
+// ProtonMail Bridge is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// ProtonMail Bridge is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with ProtonMail Bridge.  If not, see <https://www.gnu.org/licenses/>.
+
+package transfer
+
+import (
+	"sort"
+
+	"github.com/ProtonMail/gopenpgp/crypto"
+	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
+	"github.com/pkg/errors"
+)
+
+// PMAPIProvider implements import and export to/from ProtonMail server.
+type PMAPIProvider struct {
+	clientManager ClientManager
+	userID        string
+	addressID     string
+	keyRing       *crypto.KeyRing
+}
+
+// NewPMAPIProvider returns new PMAPIProvider.
+func NewPMAPIProvider(clientManager ClientManager, userID, addressID string) (*PMAPIProvider, error) {
+	keyRing, err := clientManager.GetClient(userID).KeyRingForAddressID(addressID)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to get key ring")
+	}
+
+	return &PMAPIProvider{
+		clientManager: clientManager,
+		userID:        userID,
+		addressID:     addressID,
+		keyRing:       keyRing,
+	}, nil
+}
+
+func (p *PMAPIProvider) client() pmapi.Client {
+	return p.clientManager.GetClient(p.userID)
+}
+
+// ID returns identifier of current setup of PMAPI provider.
+// Identification is unique per user.
+func (p *PMAPIProvider) ID() string {
+	return p.userID
+}
+
+// Mailboxes returns all available labels in ProtonMail account.
+func (p *PMAPIProvider) Mailboxes(includeEmpty, includeAllMail bool) ([]Mailbox, error) {
+	labels, err := p.client().ListLabels()
+	if err != nil {
+		return nil, err
+	}
+	sortedLabels := byFoldersLabels(labels)
+	sort.Sort(sortedLabels)
+
+	emptyLabelsMap := map[string]bool{}
+	if !includeEmpty {
+		messagesCounts, err := p.client().CountMessages(p.addressID)
+		if err != nil {
+			return nil, err
+		}
+		for _, messagesCount := range messagesCounts {
+			if messagesCount.Total == 0 {
+				emptyLabelsMap[messagesCount.LabelID] = true
+			}
+		}
+	}
+
+	mailboxes := getSystemMailboxes(includeAllMail)
+	for _, label := range sortedLabels {
+		if !includeEmpty && emptyLabelsMap[label.ID] {
+			continue
+		}
+
+		mailboxes = append(mailboxes, Mailbox{
+			ID:          label.ID,
+			Name:        label.Name,
+			Color:       label.Color,
+			IsExclusive: label.Exclusive == 1,
+		})
+	}
+	return mailboxes, nil
+}
+
+func getSystemMailboxes(includeAllMail bool) []Mailbox {
+	mailboxes := []Mailbox{
+		{ID: pmapi.InboxLabel, Name: "Inbox", IsExclusive: true},
+		{ID: pmapi.DraftLabel, Name: "Drafts", IsExclusive: true},
+		{ID: pmapi.SentLabel, Name: "Sent", IsExclusive: true},
+		{ID: pmapi.StarredLabel, Name: "Starred", IsExclusive: true},
+		{ID: pmapi.ArchiveLabel, Name: "Archive", IsExclusive: true},
+		{ID: pmapi.SpamLabel, Name: "Spam", IsExclusive: true},
+		{ID: pmapi.TrashLabel, Name: "Trash", IsExclusive: true},
+	}
+
+	if includeAllMail {
+		mailboxes = append(mailboxes, Mailbox{
+			ID:          pmapi.AllMailLabel,
+			Name:        "All Mail",
+			IsExclusive: true,
+		})
+	}
+
+	return mailboxes
+}
+
+type byFoldersLabels []*pmapi.Label
+
+func (l byFoldersLabels) Len() int {
+	return len(l)
+}
+
+func (l byFoldersLabels) Swap(i, j int) {
+	l[i], l[j] = l[j], l[i]
+}
+
+// Less sorts first folders, then labels, by user order.
+func (l byFoldersLabels) Less(i, j int) bool {
+	if l[i].Exclusive == 1 && l[j].Exclusive == 0 {
+		return true
+	}
+	if l[i].Exclusive == 0 && l[j].Exclusive == 1 {
+		return false
+	}
+	return l[i].Order < l[j].Order
+}