peroxide/test/liveapi/cleanup.go

// 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 custom 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, pmapi.AllMailLabel, pmapi.DraftLabel} {
		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
}