proton-bridge/internal/user/debug.go

// Copyright (c) 2023 Proton AG
//
// This file is part of Proton Mail Bridge.
//
// Proton Mail 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.
//
// Proton Mail 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 Proton Mail Bridge.  If not, see <https://www.gnu.org/licenses/>.

package user

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"mime"
	"os"
	"path/filepath"
	"strings"

	"github.com/ProtonMail/gluon/imap"
	"github.com/ProtonMail/gluon/rfc822"
	"github.com/ProtonMail/go-proton-api"
	"github.com/ProtonMail/gopenpgp/v2/constants"
	"github.com/ProtonMail/gopenpgp/v2/crypto"
	"github.com/ProtonMail/proton-bridge/v3/internal/safe"
	"github.com/ProtonMail/proton-bridge/v3/internal/usertypes"
	"github.com/ProtonMail/proton-bridge/v3/internal/vault"
	"github.com/bradenaw/juniper/xmaps"
	"github.com/bradenaw/juniper/xslices"
	"github.com/emersion/go-message"
	"github.com/sirupsen/logrus"
)

type DiagnosticMetadata struct {
	MessageIDs       []string
	Metadata         []proton.MessageMetadata
	FailedMessageIDs xmaps.Set[string]
}

type AccountMailboxMap map[string][]DiagMailboxMessage

type DiagMailboxMessage struct {
	AddressID string
	UserID    string
	ID        string
	Flags     imap.FlagSet
}

func (apm DiagnosticMetadata) BuildMailboxToMessageMap(user *User) (map[string]AccountMailboxMap, error) {
	return safe.RLockRetErr(func() (map[string]AccountMailboxMap, error) {
		result := make(map[string]AccountMailboxMap)

		mode := user.GetAddressMode()
		primaryAddrID, err := usertypes.GetPrimaryAddr(user.apiAddrs)
		if err != nil {
			return nil, fmt.Errorf("failed to get primary addr for user: %w", err)
		}

		getAccount := func(addrID string) (AccountMailboxMap, bool) {
			if mode == vault.CombinedMode {
				addrID = primaryAddrID.ID
			}

			addr := user.apiAddrs[addrID]
			if addr.Status != proton.AddressStatusEnabled {
				return nil, false
			}

			v, ok := result[addr.Email]
			if !ok {
				result[addr.Email] = make(AccountMailboxMap)
				v = result[addr.Email]
			}

			return v, true
		}

		for _, metadata := range apm.Metadata {
			for _, label := range metadata.LabelIDs {
				details, ok := user.apiLabels[label]
				if !ok {
					logrus.Warnf("User %v has message with unknown label '%v'", user.Name(), label)
					continue
				}

				if !wantLabel(details) {
					continue
				}

				account, enabled := getAccount(metadata.AddressID)
				if !enabled {
					continue
				}

				var mboxName string
				if details.Type == proton.LabelTypeSystem {
					mboxName = details.Name
				} else {
					mboxName = strings.Join(getMailboxName(details), "/")
				}

				mboxMessage := DiagMailboxMessage{
					UserID:    user.ID(),
					ID:        metadata.ID,
					AddressID: metadata.AddressID,
					Flags:     buildFlagSetFromMessageMetadata(metadata),
				}

				if v, ok := account[mboxName]; ok {
					account[mboxName] = append(v, mboxMessage)
				} else {
					account[mboxName] = []DiagMailboxMessage{mboxMessage}
				}
			}
		}
		return result, nil
	}, user.apiAddrsLock, user.apiLabelsLock)
}

func (user *User) GetDiagnosticMetadata(ctx context.Context) (DiagnosticMetadata, error) {
	failedMessages := xmaps.SetFromSlice(user.vault.SyncStatus().FailedMessageIDs)

	messageIDs, err := user.client.GetMessageIDs(ctx, "")
	if err != nil {
		return DiagnosticMetadata{}, err
	}

	meta := make([]proton.MessageMetadata, 0, len(messageIDs))

	for _, m := range xslices.Chunk(messageIDs, 100) {
		metadata, err := user.client.GetMessageMetadataPage(ctx, 0, len(m), proton.MessageFilter{ID: m})
		if err != nil {
			return DiagnosticMetadata{}, err
		}

		meta = append(meta, metadata...)
	}

	return DiagnosticMetadata{
		MessageIDs:       messageIDs,
		Metadata:         meta,
		FailedMessageIDs: failedMessages,
	}, nil
}

func (user *User) DebugDownloadMessages(
	ctx context.Context,
	path string,
	msgs map[string]DiagMailboxMessage,
	progressCB func(string, int, int),
) error {
	var err error
	safe.RLock(func() {
		err = func() error {
			total := len(msgs)
			userID := user.ID()

			counter := 1
			for _, msg := range msgs {
				if progressCB != nil {
					progressCB(userID, counter, total)
					counter++
				}

				msgDir := filepath.Join(path, msg.ID)
				if err := os.MkdirAll(msgDir, 0o700); err != nil {
					return fmt.Errorf("failed to create directory '%v':%w", msgDir, err)
				}

				message, err := user.client.GetFullMessage(ctx, msg.ID, usertypes.NewProtonAPIScheduler(user.panicHandler), proton.NewDefaultAttachmentAllocator())
				if err != nil {
					return fmt.Errorf("failed to download message '%v':%w", msg.ID, err)
				}

				if err := writeMetadata(msgDir, message.Message); err != nil {
					return err
				}

				if err := usertypes.WithAddrKR(user.apiUser, user.apiAddrs[msg.AddressID], user.vault.KeyPass(), func(_, addrKR *crypto.KeyRing) error {
					switch {
					case len(message.Attachments) > 0:
						return decodeMultipartMessage(msgDir, addrKR, message.Message, message.AttData)

					case message.MIMEType == "multipart/mixed":
						return decodePGPMessage(msgDir, addrKR, message.Message)

					default:
						return decodeSimpleMessage(msgDir, addrKR, message.Message)
					}
				}); err != nil {
					return err
				}
			}
			return nil
		}()
	}, user.apiAddrsLock, user.apiUserLock)
	return err
}

func getBodyName(path string) string {
	return filepath.Join(path, "body.txt")
}

func getBodyNameFailed(path string) string {
	return filepath.Join(path, "body_failed.txt")
}

func getBodyNamePGP(path string) string {
	return filepath.Join(path, "body.pgp")
}

func getMetadataPath(path string) string {
	return filepath.Join(path, "metadata.json")
}

func getAttachmentPathSuccess(path, id, name string) string {
	return filepath.Join(path, fmt.Sprintf("attachment_%v_%v", id, name))
}

func getAttachmentPathFailure(path, id string) string {
	return filepath.Join(path, fmt.Sprintf("attachment_%v_failed.pgp", id))
}

func decodeMultipartMessage(outPath string, kr *crypto.KeyRing, msg proton.Message, attData [][]byte) error {
	for idx, attachment := range msg.Attachments {
		if err := decodeAttachment(outPath, kr, attachment, attData[idx]); err != nil {
			return fmt.Errorf("failed to decode attachment %v of message %v: %w", attachment.ID, msg.ID, err)
		}
	}

	return decodeSimpleMessage(outPath, kr, msg)
}

func decodePGPMessage(outPath string, kr *crypto.KeyRing, msg proton.Message) error {
	var decrypted bytes.Buffer
	decrypted.Grow(len(msg.Body))

	if err := msg.DecryptInto(kr, &decrypted); err != nil {
		logrus.Warnf("Failed to decrypt pgp message %v, storing as is: %v", msg.ID, err)
		bodyPath := getBodyNamePGP(outPath)
		if err := os.WriteFile(bodyPath, []byte(msg.Body), 0o600); err != nil {
			return fmt.Errorf("failed to write pgp body to '%v': %w", bodyPath, err)
		}
	}

	bodyPath := getBodyName(outPath)

	if err := os.WriteFile(bodyPath, decrypted.Bytes(), 0o600); err != nil {
		return fmt.Errorf("failed to write pgp body to '%v': %w", bodyPath, err)
	}

	return nil
}

func decodeSimpleMessage(outPath string, kr *crypto.KeyRing, msg proton.Message) error {
	var decrypted bytes.Buffer
	decrypted.Grow(len(msg.Body))

	if err := msg.DecryptInto(kr, &decrypted); err != nil {
		logrus.Warnf("Failed to decrypt simple message %v, will try again as attachment : %v", msg.ID, err)
		return writeCustomTextPart(getBodyNameFailed(outPath), msg, err)
	}

	bodyPath := getBodyName(outPath)

	if err := os.WriteFile(bodyPath, decrypted.Bytes(), 0o600); err != nil {
		return fmt.Errorf("failed to write simple body to '%v': %w", bodyPath, err)
	}

	return nil
}

func writeMetadata(outPath string, msg proton.Message) error {
	type CustomMetadata struct {
		proton.MessageMetadata
		Header        string
		ParsedHeaders proton.Headers
		MIMEType      rfc822.MIMEType
		Attachments   []proton.Attachment
	}

	metadata := CustomMetadata{
		MessageMetadata: msg.MessageMetadata,
		Header:          msg.Header,
		ParsedHeaders:   msg.ParsedHeaders,
		MIMEType:        msg.MIMEType,
		Attachments:     msg.Attachments,
	}

	j, err := json.MarshalIndent(metadata, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to encode json for message %v: %w", msg.ID, err)
	}

	metaPath := getMetadataPath(outPath)

	if err := os.WriteFile(metaPath, j, 0o600); err != nil {
		return fmt.Errorf("failed to write metadata to '%v': %w", metaPath, err)
	}

	return nil
}

func decodeAttachment(outPath string, kr *crypto.KeyRing,
	att proton.Attachment,
	attData []byte) error {
	kps, err := base64.StdEncoding.DecodeString(att.KeyPackets)
	if err != nil {
		return err
	}

	// Use io.Multi
	attachmentReader := io.MultiReader(bytes.NewReader(kps), bytes.NewReader(attData))

	stream, err := kr.DecryptStream(attachmentReader, nil, crypto.GetUnixTime())
	if err != nil {
		logrus.
			WithField("attID", att.ID).
			WithError(err).
			Warn("Attachment decryption failed - construct")

		var pgpMessageBuffer bytes.Buffer
		pgpMessageBuffer.Grow(len(kps) + len(attData))
		pgpMessageBuffer.Write(kps)
		pgpMessageBuffer.Write(attData)

		return writeCustomAttachmentPart(getAttachmentPathFailure(outPath, att.ID), att, &crypto.PGPMessage{Data: pgpMessageBuffer.Bytes()}, err)
	}

	var decryptBuffer bytes.Buffer
	decryptBuffer.Grow(len(kps) + len(attData))

	if _, err := decryptBuffer.ReadFrom(stream); err != nil {
		logrus.
			WithField("attID", att.ID).
			WithError(err).
			Warn("Attachment decryption failed - stream")

		var pgpMessageBuffer bytes.Buffer
		pgpMessageBuffer.Grow(len(kps) + len(attData))
		pgpMessageBuffer.Write(kps)
		pgpMessageBuffer.Write(attData)

		return writeCustomAttachmentPart(getAttachmentPathFailure(outPath, att.ID), att, &crypto.PGPMessage{Data: pgpMessageBuffer.Bytes()}, err)
	}

	attachmentPath := getAttachmentPathSuccess(outPath, att.ID, att.Name)

	if err := os.WriteFile(attachmentPath, decryptBuffer.Bytes(), 0o600); err != nil {
		return fmt.Errorf("failed to write attachment %v to '%v': %w", att.ID, attachmentPath, err)
	}

	return nil
}

func writeCustomTextPart(
	outPath string,
	msg proton.Message,
	decError error,
) error {
	enc, err := crypto.NewPGPMessageFromArmored(msg.Body)
	if err != nil {
		return err
	}

	arm, err := enc.GetArmoredWithCustomHeaders(
		fmt.Sprintf("This message could not be decrypted: %v", decError),
		constants.ArmorHeaderVersion,
	)
	if err != nil {
		return err
	}

	if err := os.WriteFile(outPath, []byte(arm), 0o600); err != nil {
		return fmt.Errorf("failed to write custom message %v data to '%v': %w", msg.ID, outPath, err)
	}

	return nil
}

// writeCustomAttachmentPart writes an armored-PGP data part for an attachment that couldn't be decrypted.
func writeCustomAttachmentPart(
	outPath string,
	att proton.Attachment,
	msg *crypto.PGPMessage,
	decError error,
) error {
	arm, err := msg.GetArmoredWithCustomHeaders(
		fmt.Sprintf("This attachment could not be decrypted: %v", decError),
		constants.ArmorHeaderVersion,
	)
	if err != nil {
		return err
	}

	filename := mime.QEncoding.Encode("utf-8", att.Name+".pgp")

	var hdr message.Header

	hdr.SetContentType("application/octet-stream", map[string]string{"name": filename})
	hdr.SetContentDisposition(string(att.Disposition), map[string]string{"filename": filename})

	if err := os.WriteFile(outPath, []byte(arm), 0o600); err != nil {
		return fmt.Errorf("failed to write custom attachment %v part  to '%v': %w", att.ID, outPath, err)
	}

	return nil
}