mailgosend/internal/storage/storage.go

// Package storage provides encrypted message persistence for both SQLite (db)
// and filesystem (fs) backends. All body and raw data is encrypted with
// AES-256-GCM before being written.
package storage

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"fmt"
	"io"
	"mime"
	"mime/multipart"
	"mime/quotedprintable"
	"net/mail"
	"os"
	"path/filepath"
	"strings"
	"time"

	"ghb.freebede.com/nahakubuilder/mailgosend/internal/crypto"
	"ghb.freebede.com/nahakubuilder/mailgosend/internal/db"
	"ghb.freebede.com/nahakubuilder/mailgosend/internal/models"
)

// Store is the encrypted message store.
type Store struct {
	db      *db.DB
	crypt   *crypto.Crypto
	backend string // "db" | "fs"
	fsPath  string
}

// IncomingMessage is the raw inbound message plus parsed envelope fields.
type IncomingMessage struct {
	Raw       []byte // full RFC822
	FromEmail string
	FromName  string
	ToList    string
	CCList    string
	BCCList   string
	ReplyTo   string
	Subject   string
	Date      time.Time
	MessageID string
	SpamScore int
}

type parsedAttachment struct {
	Filename    string
	ContentType string
	Data        []byte
	ContentID   string
	Inline      bool
	MIMEPath    string
}

// New validates the backend and returns a ready Store.
func New(database *db.DB, crypt *crypto.Crypto, backend, fsPath string) (*Store, error) {
	if backend != "db" && backend != "fs" {
		return nil, fmt.Errorf("storage: unknown backend %q (want \"db\" or \"fs\")", backend)
	}
	if backend == "fs" {
		if fsPath == "" {
			return nil, fmt.Errorf("storage: fsPath required for fs backend")
		}
		if err := os.MkdirAll(fsPath, 0700); err != nil {
			return nil, fmt.Errorf("storage: create fs dir: %w", err)
		}
	}
	return &Store{db: database, crypt: crypt, backend: backend, fsPath: fsPath}, nil
}

// SaveIncoming encrypts and persists an incoming message plus its attachments.
// Returns the new message row ID.
func (s *Store) SaveIncoming(ctx context.Context, userID, mailboxID int64, msg *IncomingMessage) (int64, error) {
	uid, err := s.db.NextUID(ctx, mailboxID)
	if err != nil {
		return 0, fmt.Errorf("storage: next uid: %w", err)
	}

	bodyText, bodyHTML, attachments := parseMIME(msg.Raw)

	key, err := s.crypt.DeriveKey("messages", userID)
	if err != nil {
		return 0, fmt.Errorf("storage: derive key: %w", err)
	}

	bodyTextEnc, err := crypto.Encrypt(key, []byte(bodyText))
	if err != nil {
		return 0, fmt.Errorf("storage: encrypt body text: %w", err)
	}

	bodyHTMLEnc, err := crypto.Encrypt(key, []byte(bodyHTML))
	if err != nil {
		return 0, fmt.Errorf("storage: encrypt body html: %w", err)
	}

	rawEnc, err := crypto.Encrypt(key, msg.Raw)
	if err != nil {
		return 0, fmt.Errorf("storage: encrypt raw: %w", err)
	}

	insert := &db.MessageInsert{
		MailboxID:     mailboxID,
		UID:           uid,
		MessageID:     msg.MessageID,
		Subject:       msg.Subject,
		FromEmail:     msg.FromEmail,
		FromName:      msg.FromName,
		ToList:        msg.ToList,
		CCList:        msg.CCList,
		BCCList:       msg.BCCList,
		ReplyTo:       msg.ReplyTo,
		Date:          msg.Date,
		BodyTextEnc:   bodyTextEnc,
		BodyHTMLEnc:   bodyHTMLEnc,
		RawEnc:        rawEnc,
		SizeBytes:     int64(len(msg.Raw)),
		HasAttachment: len(attachments) > 0,
		IsRead:        false,
		IsStarred:     false,
		IsDraft:       false,
		Flags:         "",
		SpamScore:     msg.SpamScore,
	}

	messageID, err := s.db.InsertMessage(ctx, insert)
	if err != nil {
		return 0, fmt.Errorf("storage: insert message: %w", err)
	}

	if err := s.db.UpdateUsedBytes(ctx, userID, int64(len(msg.Raw))); err != nil {
		return 0, fmt.Errorf("storage: update used bytes: %w", err)
	}

	for _, att := range attachments {
		attEnc, err := crypto.Encrypt(key, att.Data)
		if err != nil {
			return 0, fmt.Errorf("storage: encrypt attachment: %w", err)
		}

		var dataPath string
		var dataEnc []byte

		if s.backend == "fs" {
			h := sha256.Sum256(att.Data)
			name := hex.EncodeToString(h[:]) + ".att"
			p := filepath.Join(s.fsPath, name)
			if err := os.WriteFile(p, attEnc, 0600); err != nil {
				return 0, fmt.Errorf("storage: write attachment file: %w", err)
			}
			dataPath = p
		} else {
			dataEnc = attEnc
		}

		attInsert := &db.AttachmentInsert{
			MessageID:   messageID,
			Filename:    att.Filename,
			ContentType: att.ContentType,
			SizeBytes:   int64(len(att.Data)),
			DataEnc:     dataEnc,
			DataPath:    dataPath,
			ContentID:   att.ContentID,
			Inline:      att.Inline,
			MIMEPath:    att.MIMEPath,
		}
		if _, err := s.db.InsertAttachment(ctx, attInsert); err != nil {
			return 0, fmt.Errorf("storage: insert attachment: %w", err)
		}
	}

	return messageID, nil
}

// GetRaw returns the decrypted RFC822 blob for a message.
func (s *Store) GetRaw(ctx context.Context, userID, messageID int64) ([]byte, error) {
	rawEnc, err := s.db.GetMessageRaw(ctx, messageID)
	if err != nil {
		return nil, fmt.Errorf("storage: get raw enc: %w", err)
	}
	if rawEnc == nil {
		return nil, fmt.Errorf("storage: message %d not found", messageID)
	}

	key, err := s.crypt.DeriveKey("messages", userID)
	if err != nil {
		return nil, fmt.Errorf("storage: derive key: %w", err)
	}

	plain, err := crypto.Decrypt(key, rawEnc)
	if err != nil {
		return nil, fmt.Errorf("storage: decrypt raw: %w", err)
	}
	return plain, nil
}

// BodyParts holds decoded body content returned by GetBodyParts.
type BodyParts struct {
	Text        string
	HTML        string
	Attachments []AttachmentMeta
}

// AttachmentMeta describes an attachment without loading its bytes.
type AttachmentMeta struct {
	Filename    string
	ContentType string
	ContentID   string // for inline images
	Inline      bool
}

// GetBodyParts decrypts a message and returns the text/HTML body and attachment list.
func (s *Store) GetBodyParts(ctx context.Context, userID, messageID int64) (*BodyParts, error) {
	raw, err := s.GetRaw(ctx, userID, messageID)
	if err != nil {
		return nil, err
	}
	text, html, atts := parseMIME(raw)
	bp := &BodyParts{Text: text, HTML: html}
	for _, a := range atts {
		bp.Attachments = append(bp.Attachments, AttachmentMeta{
			Filename:    a.Filename,
			ContentType: a.ContentType,
			ContentID:   a.ContentID,
			Inline:      a.Inline,
		})
	}
	return bp, nil
}

// GetAttachmentData decrypts and returns the bytes of the n-th attachment (0-based) for a message.
// Also returns filename and content-type. Returns an error when the attachment does not exist.
func (s *Store) GetAttachmentData(ctx context.Context, userID, messageID int64, n int) (data []byte, filename, contentType string, err error) {
	att, err := s.db.GetAttachmentByIndex(ctx, messageID, n)
	if err != nil {
		return nil, "", "", fmt.Errorf("storage: get attachment index: %w", err)
	}
	if att == nil {
		return nil, "", "", fmt.Errorf("storage: attachment %d not found for message %d", n, messageID)
	}

	key, err := s.crypt.DeriveKey("messages", userID)
	if err != nil {
		return nil, "", "", fmt.Errorf("storage: derive key: %w", err)
	}

	var encData []byte
	if att.DataPath != "" {
		encData, err = os.ReadFile(att.DataPath)
		if err != nil {
			return nil, "", "", fmt.Errorf("storage: read attachment file: %w", err)
		}
	} else {
		encData = att.DataEnc
	}
	if len(encData) == 0 {
		return nil, "", "", fmt.Errorf("storage: attachment data empty")
	}

	plain, err := crypto.Decrypt(key, encData)
	if err != nil {
		return nil, "", "", fmt.Errorf("storage: decrypt attachment: %w", err)
	}
	return plain, att.Filename, att.ContentType, nil
}

// parseMIME walks a raw RFC822 message and extracts body parts and attachments.
func parseMIME(raw []byte) (bodyText, bodyHTML string, attachments []parsedAttachment) {
	m, err := mail.ReadMessage(bytes.NewReader(raw))
	if err != nil {
		return "", "", nil
	}
	ct := m.Header.Get("Content-Type")
	body, err := io.ReadAll(m.Body)
	if err != nil {
		return "", "", nil
	}
	walkPart(ct, m.Header.Get("Content-Transfer-Encoding"),
		m.Header.Get("Content-Disposition"), m.Header.Get("Content-ID"),
		body, "1", &bodyText, &bodyHTML, &attachments)
	return bodyText, bodyHTML, attachments
}

// walkPart recursively processes one MIME part.
func walkPart(
	ctHeader, cteHeader, cdHeader, cidHeader string,
	data []byte,
	mimePath string,
	bodyText, bodyHTML *string,
	attachments *[]parsedAttachment,
) {
	mediaType, params, err := mime.ParseMediaType(ctHeader)
	if err != nil {
		// Treat unparseable as text/plain.
		mediaType = "text/plain"
		params = map[string]string{}
	}

	// Decode transfer encoding first if this is a leaf part.
	if !strings.HasPrefix(mediaType, "multipart/") {
		data = decodeCTE(cteHeader, data)
	}

	switch {
	case mediaType == "text/plain" && !isAttachment(cdHeader):
		*bodyText = string(data)

	case mediaType == "text/html" && !isAttachment(cdHeader):
		*bodyHTML = string(data)

	case strings.HasPrefix(mediaType, "multipart/"):
		boundary, ok := params["boundary"]
		if !ok {
			return
		}
		mr := multipart.NewReader(bytes.NewReader(data), boundary)
		partIdx := 1
		for {
			part, err := mr.NextPart()
			if err != nil {
				break
			}
			partData, err := io.ReadAll(part)
			if err != nil {
				part.Close()
				break
			}
			part.Close()

			subCT := part.Header.Get("Content-Type")
			if subCT == "" {
				subCT = "text/plain"
			}
			subCTE := part.Header.Get("Content-Transfer-Encoding")
			subCD := part.Header.Get("Content-Disposition")
			subCID := part.Header.Get("Content-ID")
			subPath := fmt.Sprintf("%s.%d", mimePath, partIdx)

			walkPart(subCT, subCTE, subCD, subCID, partData, subPath,
				bodyText, bodyHTML, attachments)
			partIdx++
		}

	default:
		// Treat as attachment.
		filename := filenameFrom(cdHeader, ctHeader)
		inline := strings.HasPrefix(strings.ToLower(cdHeader), "inline")
		cleanCID := strings.Trim(cidHeader, "<>")
		*attachments = append(*attachments, parsedAttachment{
			Filename:    filename,
			ContentType: mediaType,
			Data:        data,
			ContentID:   cleanCID,
			Inline:      inline,
			MIMEPath:    mimePath,
		})
	}
}

// decodeCTE applies quoted-printable or base64 decoding based on the
// Content-Transfer-Encoding header value.
func decodeCTE(cte string, data []byte) []byte {
	switch strings.ToLower(strings.TrimSpace(cte)) {
	case "quoted-printable":
		decoded, err := io.ReadAll(quotedprintable.NewReader(bytes.NewReader(data)))
		if err != nil {
			return data
		}
		return decoded
	case "base64":
		// Strip whitespace — base64 in email is line-wrapped.
		clean := strings.Map(func(r rune) rune {
			if r == '\r' || r == '\n' || r == ' ' || r == '\t' {
				return -1
			}
			return r
		}, string(data))
		decoded, err := base64.StdEncoding.DecodeString(clean)
		if err != nil {
			// Try raw/URL encoding as fallback.
			decoded, err = base64.RawStdEncoding.DecodeString(clean)
			if err != nil {
				return data
			}
		}
		return decoded
	default:
		return data
	}
}

// isAttachment returns true when the Content-Disposition header signals attachment.
func isAttachment(cd string) bool {
	if cd == "" {
		return false
	}
	disp, _, _ := mime.ParseMediaType(cd)
	return strings.EqualFold(disp, "attachment")
}

// filenameFrom extracts a filename from Content-Disposition, falling back to
// the name= param of Content-Type.
func filenameFrom(cd, ct string) string {
	if cd != "" {
		_, params, err := mime.ParseMediaType(cd)
		if err == nil {
			if name, ok := params["filename"]; ok && name != "" {
				return filepath.Base(name)
			}
		}
	}
	if ct != "" {
		_, params, err := mime.ParseMediaType(ct)
		if err == nil {
			if name, ok := params["name"]; ok && name != "" {
				return filepath.Base(name)
			}
		}
	}
	return "attachment"
}

// Ensure models import is used (Subject field type reference avoids import pruning).
var _ = models.MailboxInbox