mailgoserver/internal/smtpserver/headers.go

package smtpserver

import (
	"fmt"
	"mime"
	"net/mail"
	"strings"
	"time"

	"mailgoserver/internal/toolbox"
)

// extractMessageID scans the raw content's header block for an existing Message-ID
// header and, if its hostname doesn't match heloHostname, rewrites it to use
// heloHostname — mirroring the pre-scan in smtp_handler.handle_DATA. Unlike the Python
// version, a missing "@" or missing header entirely is handled explicitly instead of
// crashing (the approved bug fix), by falling back to a freshly generated Message-ID.
func extractMessageID(content, heloHostname string) string {
	for _, line := range strings.Split(content, "\n") {
		line = strings.TrimRight(line, "\r")
		if line == "" {
			break // end of header block
		}
		lower := strings.ToLower(line)
		if !strings.HasPrefix(lower, "message-id:") {
			continue
		}
		value := strings.TrimSpace(line[len("message-id:"):])
		value = strings.Trim(value, "<>")
		at := strings.LastIndex(value, "@")
		if at < 0 {
			break // malformed header, no "@" — fall through to generating a fresh one
		}
		prefix, hostname := value[:at], value[at+1:]
		if !strings.EqualFold(hostname, heloHostname) {
			return fmt.Sprintf("%s@%s", prefix, heloHostname)
		}
		return value
	}
	return toolbox.GenerateMessageID(heloHostname)
}

// existingHeaders parses the raw header block into a lowercase-keyed map of the first
// value seen per header name, folding continuation lines, mirroring the case-insensitive
// existing-header lookups in _ensure_required_headers.
func existingHeaders(content string) map[string]string {
	lines := strings.Split(content, "\n")
	out := map[string]string{}
	var lastKey string
	for _, raw := range lines {
		line := strings.TrimRight(raw, "\r")
		if line == "" {
			break
		}
		if (strings.HasPrefix(line, " ") || strings.HasPrefix(line, "\t")) && lastKey != "" {
			out[lastKey] += " " + strings.TrimSpace(line)
			continue
		}
		idx := strings.Index(line, ":")
		if idx < 0 {
			continue
		}
		key := strings.ToLower(strings.TrimSpace(line[:idx]))
		val := strings.TrimSpace(line[idx+1:])
		out[key] = val
		lastKey = key
	}
	return out
}

// splitHeadersBody separates the header block from the body on the first blank line,
// accepting either CRLF or bare-LF line endings (source content is LF-only from the
// SMTP DATA decode; ensureRequiredHeaders' own output is CRLF).
func splitHeadersBody(content string) (headerBlock, body string) {
	if idx := strings.Index(content, "\r\n\r\n"); idx >= 0 {
		return content[:idx], content[idx+4:]
	}
	if idx := strings.Index(content, "\n\n"); idx >= 0 {
		return content[:idx], content[idx+2:]
	}
	return content, ""
}

// ensureRequiredHeaders performs a full header-block *replacement* (not augmentation),
// mirroring smtp_handler._ensure_required_headers exactly: a fixed, ordered whitelist
// of headers is emitted, copying values from the original message where present and
// defaulting where absent; anything outside that whitelist is dropped, then the
// domain's custom headers plus X-Originating-IP/X-Mailer/X-Priority are appended (only
// if not already present under the same name).
func ensureRequiredHeaders(content, messageID string, envelopeRcptTos []string, mailFrom string, customHeaders [][2]string) string {
	headerBlock, body := splitHeadersBody(content)
	existing := existingHeaders(headerBlock)

	var out []string
	out = append(out, "Message-ID: <"+messageID+">")

	if v, ok := existing["date"]; ok {
		out = append(out, "Date: "+v)
	} else {
		out = append(out, "Date: "+time.Now().Format(time.RFC1123Z))
	}

	if v, ok := existing["mime-version"]; ok {
		out = append(out, "MIME-Version: "+v)
	} else {
		out = append(out, "MIME-Version: 1.0")
	}

	if v, ok := existing["to"]; ok {
		out = append(out, "To: "+v)
	} else {
		out = append(out, "To: "+strings.Join(envelopeRcptTos, ", "))
	}

	if v, ok := existing["cc"]; ok {
		out = append(out, "Cc: "+v)
	}

	if v, ok := existing["from"]; ok {
		out = append(out, "From: "+v)
	} else {
		out = append(out, "From: "+mailFrom)
	}

	if v, ok := existing["subject"]; ok {
		out = append(out, "Subject: "+v)
	} else {
		out = append(out, "Subject: ")
	}

	if v, ok := existing["content-type"]; ok {
		out = append(out, "Content-Type: "+v)
	} else {
		out = append(out, `Content-Type: text/plain; charset=UTF-8; format=flowed`)
	}

	if v, ok := existing["content-transfer-encoding"]; ok {
		out = append(out, "Content-Transfer-Encoding: "+v)
	} else {
		out = append(out, "Content-Transfer-Encoding: 7bit")
	}

	for _, kv := range customHeaders {
		if _, already := existing[strings.ToLower(kv[0])]; already {
			continue
		}
		out = append(out, kv[0]+": "+kv[1])
	}

	return strings.Join(out, "\r\n") + "\r\n\r\n" + body
}

// getContentType mirrors smtp_handler.get_content_type: prefer the part's own type,
// fall back to extension sniffing, then a small fixed extension map.
func getContentType(partContentType, filename string) string {
	if partContentType != "" && partContentType != "application/octet-stream" {
		return partContentType
	}
	if guessed := mime.TypeByExtension(extOf(filename)); guessed != "" {
		return guessed
	}
	switch strings.ToLower(extOf(filename)) {
	case ".txt":
		return "text/plain"
	case ".csv":
		return "text/csv"
	case ".jpg", ".jpeg":
		return "image/jpeg"
	case ".png":
		return "image/png"
	case ".gif":
		return "image/gif"
	case ".pdf":
		return "application/pdf"
	case ".json":
		return "application/json"
	case ".xml":
		return "application/xml"
	case ".html", ".htm":
		return "text/html"
	default:
		return "application/octet-stream"
	}
}

func extOf(filename string) string {
	if i := strings.LastIndex(filename, "."); i >= 0 {
		return filename[i:]
	}
	return ""
}

// parseAddressList mirrors the lowercase address parsing used to classify To/Cc/Bcc.
func parseAddressList(headerValue string) []string {
	if strings.TrimSpace(headerValue) == "" {
		return nil
	}
	addrs, err := mail.ParseAddressList(headerValue)
	if err != nil {
		return nil
	}
	out := make([]string, len(addrs))
	for i, a := range addrs {
		out[i] = strings.ToLower(a.Address)
	}
	return out
}