feat(GODT-2759): Add prompt to download missing messages for analysis

This will download the missing messages into a temporary directory and
decrypt them along with the metadata so we can attempt analyze them once
submitted to see what is going wrong.

internal/user/debug.go

@@ -18,16 +18,27 @@
 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/vault"
 	"github.com/bradenaw/juniper/xmaps"
 	"github.com/bradenaw/juniper/xslices"
+	"github.com/emersion/go-message"
 	"github.com/sirupsen/logrus"
 )
 
@@ -37,11 +48,13 @@ type DiagnosticMetadata struct {
 	FailedMessageIDs xmaps.Set[string]
 }
 
-type AccountMailboxMap map[string][]MailboxMessage
+type AccountMailboxMap map[string][]DiagMailboxMessage
 
-type MailboxMessage struct {
-	ID    string
-	Flags imap.FlagSet
+type DiagMailboxMessage struct {
+	AddressID string
+	UserID    string
+	ID        string
+	Flags     imap.FlagSet
 }
 
 func (apm DiagnosticMetadata) BuildMailboxToMessageMap(user *User) (map[string]AccountMailboxMap, error) {
@@ -97,15 +110,17 @@ func (apm DiagnosticMetadata) BuildMailboxToMessageMap(user *User) (map[string]A
 					mboxName = strings.Join(getMailboxName(details), "/")
 				}
 
-				mboxMessage := MailboxMessage{
-					ID:    metadata.ID,
-					Flags: buildFlagSetFromMessageMetadata(metadata),
+				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] = []MailboxMessage{mboxMessage}
+					account[mboxName] = []DiagMailboxMessage{mboxMessage}
 				}
 			}
 		}
@@ -138,3 +153,267 @@ func (user *User) GetDiagnosticMetadata(ctx context.Context) (DiagnosticMetadata
 		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, 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 := 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
+}