mailgoserver/internal/smtpserver/session.go

package smtpserver

import (
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/emersion/go-smtp"
	"gopkg.in/ini.v1"
	"mailgoserver/internal/abuseguard"
	"mailgoserver/internal/db"
	"mailgoserver/internal/dkim"
	"mailgoserver/internal/mailstore"
	"mailgoserver/internal/relay"
	"mailgoserver/internal/toolbox"
)

// Backend holds the shared dependencies every connection's Session uses, mirroring the
// constructor args threaded through smtp_handler.EnhancedCustomSMTPHandler /
// email_server/server_runner.py.
type Backend struct {
	DB                  *db.DB
	DKIM                *dkim.Manager
	Relay               *relay.Relay
	Cfg                 *ini.File
	Mailstore           *mailstore.Store
	Logger              *toolbox.Logger
	HeloHostname        string
	AttachmentsBasePath string
}

func (b *Backend) NewSession(c *smtp.Conn) (smtp.Session, error) {
	host, _, _ := net.SplitHostPort(c.Conn().RemoteAddr().String())
	if host == "" {
		host = c.Conn().RemoteAddr().String()
	}
	return &Session{backend: b, conn: c, peerIP: host}, nil
}

// Session implements smtp.Session + smtp.AuthSession for one SMTP connection, mirroring
// EnhancedCustomSMTPHandler's per-connection behavior in smtp_handler.py.
type Session struct {
	backend *Backend
	conn    *smtp.Conn
	peerIP  string

	authenticatedSender  *db.Sender
	authenticatedMailbox *db.Mailbox // set instead of authenticatedSender when auth used an app password
	authType             string      // "sender" | "mailbox" | "ip" | ""
	authorizedDomain     string
	username             string

	mailFrom           string
	mailFromAuthorized bool // true only via an existing authorized path (sender/IP) on one of our own domains
	rcptTos            []string
	localMailboxes     map[string]*db.Mailbox // lowercased rcpt -> resolved local mailbox, set in Rcpt
}

func (s *Session) Reset() {
	s.mailFrom = ""
	s.mailFromAuthorized = false
	s.rcptTos = nil
	s.localMailboxes = nil
}

func (s *Session) Logout() error { return nil }

// Mail mirrors EnhancedCustomSMTPHandler.handle_MAIL, delegating authorization to
// validateSenderAuthorization (== auth.validate_sender_authorization).
func (s *Session) Mail(from string, opts *smtp.MailOptions) error {
	accept, authorized, message := s.validateSenderAuthorization(from)
	if !accept {
		return &smtp.SMTPError{Code: 550, EnhancedCode: smtp.NoEnhancedCode, Message: message}
	}
	s.mailFrom = from
	s.mailFromAuthorized = authorized
	return nil
}

// validateSenderAuthorization mirrors auth.validate_sender_authorization for every
// domain configured on this server — that part is byte-for-byte unchanged: senders
// unauthorized or unverified on OUR OWN domains are still hard-rejected here, exactly
// as before, to prevent spoofing/open-relay for domains we're responsible for.
//
// One addition: when this server has local mailbox storage enabled (Mailstore != nil),
// a MAIL FROM on a domain we don't manage at all is now provisionally accepted
// (accept=true, authorized=false) instead of hard-rejected — otherwise this server
// could never receive genuine inbound mail from the internet, since every external
// sender's domain is by definition "not configured here". Rcpt enforces that a
// provisionally-accepted sender may only deliver to a local mailbox, never relay
// onward, so this cannot be used as an open relay.
func (s *Session) validateSenderAuthorization(mailFrom string) (accept, authorized bool, message string) {
	if mailFrom == "" {
		return false, false, "No sender address provided"
	}
	fromDomain := domainOfAddr(mailFrom)
	if fromDomain == "" {
		return false, false, "Invalid sender address format"
	}

	dom, err := s.backend.DB.GetDomainByName(fromDomain)
	if err != nil {
		s.backend.Logger.Error("domain lookup failed: %v", err)
		return false, false, fmt.Sprintf("Not authorized to send for domain %s", fromDomain)
	}
	if dom == nil {
		if s.backend.Mailstore != nil {
			return true, false, fmt.Sprintf("Domain %s not configured here; accepted for possible local delivery only", fromDomain)
		}
		return false, false, fmt.Sprintf("Domain %s is not configured on this server", fromDomain)
	}
	if !dom.IsVerified {
		return false, false, fmt.Sprintf("Domain %s has not completed DNS ownership verification yet", fromDomain)
	}

	if s.authenticatedSender != nil {
		sender := s.authenticatedSender
		if sender.CanSendAs(mailFrom) {
			return true, true, fmt.Sprintf("Sender authorized to send as %s", mailFrom)
		}
		_ = s.backend.DB.LogAuthAttempt("sender_validation", fmt.Sprintf("%s -> %s", sender.Email, mailFrom), s.peerIP, false, "")
		abuseguard.RecordFailureAndMaybeBlacklist(s.backend.DB, s.backend.Cfg, s.backend.Logger, s.peerIP)
		return false, false, fmt.Sprintf("Sender %s not authorized to send as %s", sender.Email, mailFrom)
	}

	// A mailbox (authenticated via app password) may send as its own primary address,
	// or as any of its active send-as-enabled aliases — never as an arbitrary address,
	// even within a domain it happens to own a mailbox on.
	if s.authenticatedMailbox != nil {
		mbox := s.authenticatedMailbox
		if strings.EqualFold(mailFrom, mbox.Email) {
			return true, true, fmt.Sprintf("Mailbox authorized to send as %s", mailFrom)
		}
		if canSendAs, err := s.backend.DB.MailboxCanSendAs(mbox.ID, mailFrom); err == nil && canSendAs {
			return true, true, fmt.Sprintf("Mailbox authorized to send as alias %s", mailFrom)
		}
		_ = s.backend.DB.LogAuthAttempt("mailbox_validation", fmt.Sprintf("%s -> %s", mbox.Email, mailFrom), s.peerIP, false, "")
		abuseguard.RecordFailureAndMaybeBlacklist(s.backend.DB, s.backend.Cfg, s.backend.Logger, s.peerIP)
		return false, false, fmt.Sprintf("Mailbox %s not authorized to send as %s", mbox.Email, mailFrom)
	}

	wl, err := s.backend.DB.GetWhitelistedIP(s.peerIP, fromDomain)
	if err != nil {
		s.backend.Logger.Error("IP authorization lookup failed: %v", err)
		return false, false, fmt.Sprintf("Not authorized to send for domain %s", fromDomain)
	}
	if wl != nil {
		s.authType = "ip"
		s.authorizedDomain = fromDomain
		s.username = "IP:" + s.peerIP
		_ = s.backend.DB.LogAuthAttempt("ip", fmt.Sprintf("%s -> %s", s.peerIP, fromDomain), s.peerIP, true, fmt.Sprintf("IP %s authorized for domain %s", s.peerIP, fromDomain))
		return true, true, fmt.Sprintf("IP authorized for domain %s", fromDomain)
	}
	_ = s.backend.DB.LogAuthAttempt("ip", fmt.Sprintf("%s -> %s", s.peerIP, fromDomain), s.peerIP, false, fmt.Sprintf("IP %s not authorized for domain %s", s.peerIP, fromDomain))
	abuseguard.RecordFailureAndMaybeBlacklist(s.backend.DB, s.backend.Cfg, s.backend.Logger, s.peerIP)
	return false, false, fmt.Sprintf("Not authorized to send for domain %s", fromDomain)
}

func domainOfAddr(address string) string {
	i := strings.LastIndex(address, "@")
	if i < 0 {
		return ""
	}
	return strings.ToLower(address[i+1:])
}

// Rcpt mirrors handle_RCPT for the pure-relay case (still accept-all for an authorized
// sender relaying to an external address — unchanged), and adds local-mailbox
// resolution: a recipient on one of our own configured+verified domains must resolve
// to a real mailbox, or is rejected with 550 "No such mailbox" — matching how a real
// MTA rejects unknown local recipients at RCPT time. A recipient that resolves to
// neither a local mailbox nor an authorized-to-relay sender's target is rejected with
// "Relay access denied" — the anti-open-relay invariant for provisionally-accepted
// senders (see validateSenderAuthorization).
func (s *Session) Rcpt(to string, opts *smtp.RcptOptions) error {
	domain := domainOfAddr(to)
	var localDomain *db.Domain
	if domain != "" && s.backend.Mailstore != nil {
		if dom, err := s.backend.DB.GetDomainByName(domain); err == nil && dom != nil && dom.IsVerified {
			localDomain = dom
		}
	}

	if localDomain != nil {
		mbox, err := s.backend.Mailstore.ResolveRecipient(to)
		if err != nil || mbox == nil {
			return &smtp.SMTPError{Code: 550, EnhancedCode: smtp.NoEnhancedCode, Message: "No such mailbox"}
		}
		if blocked, _ := s.backend.DB.IsBlocked(mbox.ID, s.mailFrom); blocked {
			return &smtp.SMTPError{Code: 550, EnhancedCode: smtp.NoEnhancedCode, Message: "Message rejected"}
		}
		if s.localMailboxes == nil {
			s.localMailboxes = map[string]*db.Mailbox{}
		}
		s.localMailboxes[strings.ToLower(to)] = mbox
		s.rcptTos = append(s.rcptTos, to)
		return nil
	}

	if !s.mailFromAuthorized {
		return &smtp.SMTPError{Code: 550, EnhancedCode: smtp.NoEnhancedCode, Message: "Relay access denied"}
	}
	s.rcptTos = append(s.rcptTos, to)
	return nil
}

func internalError(msg string) error {
	return &smtp.SMTPError{Code: 550, EnhancedCode: smtp.NoEnhancedCode, Message: msg}
}

// Data mirrors EnhancedCustomSMTPHandler.handle_DATA end to end: Message-ID
// extraction/rehost, full header rebuild, DKIM signing, attachment extraction/storage,
// relay delivery, and EmailLog/EmailRecipientLog/EmailAttachment persistence.
func (s *Session) Data(r io.Reader) error {
	raw, err := io.ReadAll(r)
	if err != nil {
		return internalError("Internal server error")
	}
	content := string(raw)

	messageID := extractMessageID(content, s.backend.HeloHostname)
	senderDomain := domainOfAddr(s.mailFrom)

	var customHeaders [][2]string
	if senderDomain != "" {
		customHeaders, _ = s.backend.DKIM.GetActiveCustomHeaders(senderDomain)
	}
	customHeaders = append(customHeaders,
		[2]string{"X-Originating-IP", "[" + s.peerIP + "]"},
		[2]string{"X-Mailer", "NetBro Mail Server 1.0"},
		[2]string{"X-Priority", "3"},
	)

	rebuilt := ensureRequiredHeaders(content, messageID, s.rcptTos, s.mailFrom, customHeaders)

	signedContent := rebuilt
	dkimSigned := false
	if senderDomain != "" {
		signedContent = s.backend.DKIM.Sign(rebuilt, senderDomain)
		dkimSigned = signedContent != rebuilt
	}

	rebuiltHeaders := existingHeaders(rebuilt)
	toHeader := rebuiltHeaders["to"]
	ccHeader := rebuiltHeaders["cc"]
	subject := rebuiltHeaders["subject"]

	// Attachment storage: only if the authenticated sender or whitelisted IP opted in.
	storeMessage := false
	if sender, _ := s.backend.DB.GetSenderByEmail(s.mailFrom); sender != nil && sender.StoreMessageContent {
		storeMessage = true
	} else if wl, _ := s.backend.DB.GetWhitelistedIP(s.peerIP, senderDomain); wl != nil && wl.StoreMessageContent {
		storeMessage = true
	}

	parsed, parseErr := parseMessage(raw)

	type savedAttachment struct {
		Filename, ContentType, FilePath string
		Size                            int64
	}
	var toSave []savedAttachment
	if storeMessage && parseErr == nil && len(parsed.Attachments) > 0 {
		usernameOrIP := s.username
		if usernameOrIP == "" && s.peerIP != "" {
			usernameOrIP = sanitizePathSegment(s.peerIP, ":")
		} else {
			usernameOrIP = sanitizePathSegment(usernameOrIP, "/\\")
		}
		storagePath := attachmentStoragePath(s.backend.AttachmentsBasePath, senderDomain, usernameOrIP, time.Now())
		if err := os.MkdirAll(storagePath, 0o755); err == nil {
			prefix := cleanMessageIDPrefix(messageID)
			for _, a := range parsed.Attachments {
				filename := prefix + "_" + a.Filename
				fullPath := filepath.Join(storagePath, filename)
				if err := os.WriteFile(fullPath, a.Data, 0o644); err == nil {
					toSave = append(toSave, savedAttachment{Filename: a.Filename, ContentType: a.ContentType, FilePath: fullPath, Size: int64(len(a.Data))})
				} else {
					s.backend.Logger.Error("Failed to write attachment %s: %v", filename, err)
				}
			}
		}
	}

	// Classify each envelope recipient as to/cc/bcc by presence in the To/Cc headers —
	// anything not literally present in either is inferred BCC.
	toList := parseAddressList(toHeader)
	ccList := parseAddressList(ccHeader)
	recipientTypes := make([]string, len(s.rcptTos))
	for i, rcpt := range s.rcptTos {
		lower := strings.ToLower(rcpt)
		switch {
		case containsStr(toList, lower):
			recipientTypes[i] = "to"
		case containsStr(ccList, lower):
			recipientTypes[i] = "cc"
		default:
			recipientTypes[i] = "bcc"
		}
	}

	// Split recipients resolved to a local mailbox in Rcpt from everything else
	// (still relayed exactly as before — unchanged for every non-local recipient).
	var localRcpts, localTypes, relayRcpts, relayTypes []string
	for i, rcpt := range s.rcptTos {
		if _, ok := s.localMailboxes[strings.ToLower(rcpt)]; ok {
			localRcpts = append(localRcpts, rcpt)
			localTypes = append(localTypes, recipientTypes[i])
		} else {
			relayRcpts = append(relayRcpts, rcpt)
			relayTypes = append(relayTypes, recipientTypes[i])
		}
	}

	var results []relay.Result
	if len(relayRcpts) > 0 {
		results = s.backend.Relay.RelayEmailAsync(s.mailFrom, relayRcpts, signedContent, relayTypes)
	}
	if len(localRcpts) > 0 {
		results = append(results, s.deliverLocally(localRcpts, localTypes, signedContent, messageID, subject)...)
	}

	allSucceeded := len(results) > 0
	for _, res := range results {
		if res.Status != "success" {
			allSucceeded = false
		}
	}

	var emailHeaders, messageBody string
	if parseErr == nil {
		emailHeaders = strings.Join(parsed.HeaderLines, "\n")
		messageBody = parsed.BodyText
	}

	logID, logErr := s.backend.Relay.LogEmail(s.backend.Cfg, s.peerIP, s.mailFrom, toHeader, ccHeader, "", subject, emailHeaders, messageBody, messageID, s.username, dkimSigned, results)
	if logErr != nil {
		s.backend.Logger.Error("Failed to log email: %v", logErr)
	} else {
		for _, a := range toSave {
			if err := s.backend.DB.InsertEmailAttachment(db.EmailAttachment{
				EmailLogID: logID, Filename: a.Filename, ContentType: a.ContentType, FilePath: a.FilePath, Size: a.Size,
			}); err != nil {
				s.backend.Logger.Error("Failed to record attachment %s: %v", a.Filename, err)
			}
		}
	}

	if allSucceeded {
		return &smtp.SMTPError{Code: 250, EnhancedCode: smtp.NoEnhancedCode, Message: "Message accepted for delivery"}
	}
	return &smtp.SMTPError{Code: 550, EnhancedCode: smtp.NoEnhancedCode, Message: "Message relay failed"}
}

// deliverLocally runs the inbound DKIM/SPF/spam checks once for the message (they
// don't vary per recipient at this milestone — no per-mailbox allow/block-list yet)
// and stores it into each resolved local mailbox, producing one relay.Result per
// recipient so it can be merged into the same LogEmail/allSucceeded logic as relay
// results.
func (s *Session) deliverLocally(rcpts, types []string, signedContent, messageID, subject string) []relay.Result {
	senderDomain := domainOfAddr(s.mailFrom)
	dkimPass := senderDomain != "" && dkim.VerifyInbound(signedContent, senderDomain)
	spfPass := mailstore.CheckSPF(s.mailFrom, s.peerIP)
	heuristicScore := mailstore.SpamScore(s.peerIP, map[string]string{"subject": subject}, dkimPass, spfPass)
	rejectScore := s.backend.Cfg.Section("Mailstore").Key("spam_reject_score").MustInt(5)
	rspamdEnabled := s.backend.Cfg.Section("Rspamd").Key("enabled").MustBool(false)
	rspamdURL := s.backend.Cfg.Section("Rspamd").Key("url").MustString("http://127.0.0.1:11333")
	rspamdRejectScore := s.backend.Cfg.Section("Rspamd").Key("reject_score").MustInt(15)

	results := make([]relay.Result, 0, len(rcpts))
	for i, rcpt := range rcpts {
		mbox := s.localMailboxes[strings.ToLower(rcpt)]
		folder := "INBOX"
		markRead := false

		// An explicit per-mailbox allow-list entry bypasses spam scoring entirely —
		// the built-in heuristic and optional rspamd check both run regardless of each
		// other (additive, not either/or), but neither runs at all once allow-listed.
		spamGated := false
		if allowed, _ := s.backend.DB.IsAllowed(mbox.ID, s.mailFrom); !allowed {
			quarantine := heuristicScore >= rejectScore
			hardReject := false
			if rspamdEnabled {
				if score, rAction, err := mailstore.CheckRspamd(rspamdURL, []byte(signedContent), s.mailFrom, rcpt); err == nil {
					// rspamd's own "reject" action is a considered policy decision
					// (DNSBL hit, greylisting, etc.) worth still hard-rejecting at
					// SMTP time to avoid backscatter; a bare score threshold hit
					// (from either scorer) is quarantined instead of rejected, so a
					// false positive is recoverable from the Spam folder rather than
					// silently bounced with no trace.
					if rAction == "reject" {
						hardReject = true
					} else if score >= float64(rspamdRejectScore) {
						quarantine = true
					}
				}
				// rspamd unreachable/erroring must not block mail — errors are swallowed,
				// the built-in heuristic above is still the baseline gate either way.
			}
			if hardReject {
				results = append(results, relay.Result{Recipient: rcpt, RecipientType: types[i], Status: "failed", ErrorCode: "550", ErrorMessage: "Message rejected as spam"})
				continue
			}
			if quarantine {
				folder = "Spam"
				spamGated = true
			}
		}

		// Filter rules organize legitimate mail the recipient already trusts arriving
		// in their INBOX — a quarantined message skips them entirely and always lands
		// in Spam, rather than a rule accidentally routing spam back into view.
		if !spamGated {
			action, err := s.backend.Mailstore.ApplyRules(mbox.ID, map[string]string{"from": s.mailFrom, "to": rcpt, "subject": subject})
			if err != nil {
				results = append(results, relay.Result{Recipient: rcpt, RecipientType: types[i], Status: "failed", ErrorCode: "450", ErrorMessage: err.Error()})
				continue
			}
			if action.Drop {
				results = append(results, relay.Result{Recipient: rcpt, RecipientType: types[i], Status: "success", ServerResponse: "Discarded by filter rule"})
				continue
			}
			if action.Folder != "" {
				folder = action.Folder
			}
			markRead = action.MarkRead
		}

		uid, err := s.backend.Mailstore.StoreMessage(mbox.ID, folder, []byte(signedContent), messageID, s.mailFrom, subject)
		if err != nil {
			errCode, errMsg := "450", err.Error()
			if err == mailstore.ErrQuotaExceeded {
				errCode, errMsg = "552", "Mailbox quota exceeded"
			}
			results = append(results, relay.Result{Recipient: rcpt, RecipientType: types[i], Status: "failed", ErrorCode: errCode, ErrorMessage: errMsg})
			continue
		}
		if markRead {
			if err := s.backend.DB.SetMessageFlags(mbox.ID, uid, `\Seen`); err != nil {
				s.backend.Logger.Error("mark_read rule failed to set flag for message %d: %v", uid, err)
			}
		}
		serverResponse := "Delivered to local mailbox"
		if spamGated {
			serverResponse = "Quarantined to Spam folder"
		}
		results = append(results, relay.Result{Recipient: rcpt, RecipientType: types[i], Status: "success", ServerResponse: serverResponse})
	}
	return results
}

func containsStr(list []string, s string) bool {
	for _, v := range list {
		if v == s {
			return true
		}
	}
	return false
}