ultisuite-backend/internal/mail/imap/attachments.go

package imap

import (
	"bytes"
	"encoding/base64"
	"io"
	"mime"
	"mime/multipart"
	"mime/quotedprintable"
	"net/mail"
	"strings"

	"github.com/ultisuite/ulti-backend/internal/mail/limits"
)

// AttachmentPart is a decoded MIME body part stored as an attachment or inline resource.
type AttachmentPart struct {
	Filename    string
	ContentType string
	ContentID   string
	IsInline    bool
	Data        []byte
}

// ExtractAttachments parses raw RFC 822 message bytes and returns attachment parts.
// Body text/plain and text/html parts are skipped. Non-text parts are collected when
// Content-Disposition is attachment, or inline with a filename. Parts exceeding
// limits.MaxAttachmentBytes are skipped; collection stops at limits.MaxAttachmentsPerMessage.
func ExtractAttachments(raw []byte) ([]AttachmentPart, error) {
	if len(raw) == 0 {
		return nil, nil
	}

	msg, err := mail.ReadMessage(bytes.NewReader(raw))
	if err != nil {
		return nil, err
	}

	contentType := msg.Header.Get("Content-Type")
	if contentType == "" {
		return nil, nil
	}

	mediaType, params, err := mime.ParseMediaType(contentType)
	if err != nil || !strings.HasPrefix(mediaType, "multipart/") {
		return nil, nil
	}

	return extractAttachmentsFromMultipart(msg.Body, params["boundary"]), nil
}

func extractAttachmentsFromMultipart(r io.Reader, boundary string) []AttachmentPart {
	var attachments []AttachmentPart

	mr := multipart.NewReader(r, boundary)
	for {
		part, err := mr.NextPart()
		if err != nil {
			break
		}
		if len(attachments) >= limits.MaxAttachmentsPerMessage {
			break
		}

		partType := part.Header.Get("Content-Type")
		mediaType, params, _ := mime.ParseMediaType(partType)
		if mediaType == "" {
			mediaType = "application/octet-stream"
		}

		switch {
		case mediaType == "text/plain", mediaType == "text/html":
			continue
		case strings.HasPrefix(mediaType, "multipart/"):
			nested := extractAttachmentsFromMultipart(part, params["boundary"])
			for _, att := range nested {
				if len(attachments) >= limits.MaxAttachmentsPerMessage {
					break
				}
				attachments = append(attachments, att)
			}
		default:
			if att, ok := partToAttachment(part, mediaType, params); ok {
				attachments = append(attachments, att)
			}
		}
	}

	return attachments
}

func partToAttachment(part *multipart.Part, mediaType string, typeParams map[string]string) (AttachmentPart, bool) {
	if strings.HasPrefix(mediaType, "text/") {
		return AttachmentPart{}, false
	}

	disposition, dispParams, _ := mime.ParseMediaType(part.Header.Get("Content-Disposition"))
	filename := dispParams["filename"]
	if filename == "" {
		filename = typeParams["name"]
	}

	isInline := strings.EqualFold(disposition, "inline")
	isAttachment := strings.EqualFold(disposition, "attachment")

	if !isAttachment && !(isInline && filename != "") {
		return AttachmentPart{}, false
	}

	data, err := io.ReadAll(part)
	if err != nil {
		return AttachmentPart{}, false
	}

	data, err = decodePartBody(part.Header.Get("Content-Transfer-Encoding"), data)
	if err != nil || len(data) > limits.MaxAttachmentBytes {
		return AttachmentPart{}, false
	}

	return AttachmentPart{
		Filename:    filename,
		ContentType: mediaType,
		ContentID:   normalizeContentID(part.Header.Get("Content-ID")),
		IsInline:    isInline,
		Data:        data,
	}, true
}

func normalizeContentID(raw string) string {
	return strings.Trim(raw, "<> \t")
}

func decodePartBody(transferEncoding string, data []byte) ([]byte, error) {
	switch strings.ToLower(strings.TrimSpace(transferEncoding)) {
	case "base64":
		return decodeBase64Body(data)
	case "quoted-printable":
		return io.ReadAll(quotedprintable.NewReader(bytes.NewReader(data)))
	default:
		return data, nil
	}
}

func decodeBase64Body(data []byte) ([]byte, error) {
	clean := bytes.Map(func(r rune) rune {
		switch r {
		case '\r', '\n', ' ', '\t':
			return -1
		default:
			return r
		}
	}, data)
	decoded := make([]byte, base64.StdEncoding.DecodedLen(len(clean)))
	n, err := base64.StdEncoding.Decode(decoded, clean)
	if err != nil {
		return nil, err
	}
	return decoded[:n], nil
}