We build too many walls and not enough bridges

test/liveapi/cleanup.go

@@ -0,0 +1,132 @@
+// Copyright (c) 2020 Proton Technologies AG
+//
+// This file is part of ProtonMail Bridge.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 liveapi
+
+import (
+	"time"
+
+	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
+	"github.com/pkg/errors"
+)
+
+func cleanup(client *pmapi.Client) error {
+	if err := cleanSystemFolders(client); err != nil {
+		return errors.Wrap(err, "failed to clean system folders")
+	}
+	if err := cleanCustomLables(client); err != nil {
+		return errors.Wrap(err, "failed to clean cusotm labels")
+	}
+	if err := cleanTrash(client); err != nil {
+		return errors.Wrap(err, "failed to clean trash")
+	}
+	return nil
+}
+
+func cleanSystemFolders(client *pmapi.Client) error {
+	for _, labelID := range []string{pmapi.InboxLabel, pmapi.SentLabel, pmapi.ArchiveLabel} {
+		for {
+			messages, total, err := client.ListMessages(&pmapi.MessagesFilter{
+				PageSize: 150,
+				LabelID:  labelID,
+			})
+			if err != nil {
+				return errors.Wrap(err, "failed to list messages")
+			}
+
+			if total == 0 {
+				break
+			}
+
+			messageIDs := []string{}
+			for _, message := range messages {
+				messageIDs = append(messageIDs, message.ID)
+			}
+
+			if err := client.DeleteMessages(messageIDs); err != nil {
+				return errors.Wrap(err, "failed to delete messages")
+			}
+
+			if total == len(messages) {
+				break
+			}
+		}
+	}
+	return nil
+}
+
+func cleanCustomLables(client *pmapi.Client) error {
+	labels, err := client.ListLabels()
+	if err != nil {
+		return errors.Wrap(err, "failed to list labels")
+	}
+
+	for _, label := range labels {
+		if err := emptyFolder(client, label.ID); err != nil {
+			return errors.Wrap(err, "failed to empty label")
+		}
+		if err := client.DeleteLabel(label.ID); err != nil {
+			return errors.Wrap(err, "failed to delete label")
+		}
+	}
+
+	return nil
+}
+
+func cleanTrash(client *pmapi.Client) error {
+	for {
+		_, total, err := client.ListMessages(&pmapi.MessagesFilter{
+			PageSize: 1,
+			LabelID:  pmapi.TrashLabel,
+		})
+		if err == nil && total == 0 {
+			break
+		}
+
+		err = emptyFolder(client, pmapi.TrashLabel)
+		if err == nil {
+			break
+		}
+		if err.Error() == "Folder or label is currently being emptied" {
+			time.Sleep(500 * time.Millisecond)
+			continue
+		}
+		return errors.Wrap(err, "failed to empty trash")
+	}
+	return nil
+}
+
+func emptyFolder(client *pmapi.Client, labelID string) error {
+	err := client.EmptyFolder(labelID, "")
+	if err != nil {
+		return err
+	}
+	for {
+		_, total, err := client.ListMessages(&pmapi.MessagesFilter{
+			PageSize: 1,
+			LabelID:  labelID,
+		})
+		if err != nil {
+			return err
+		}
+		if total == 0 {
+			break
+		}
+		time.Sleep(1 * time.Second)
+	}
+	return nil
+}