gomail/internal/imap/commands.go

package imap

import (
	"fmt"
	"strconv"
	"strings"

	"gomail/internal/db"
)

func (s *session) cmdCapability(tag string) {
	caps := "CAPABILITY IMAP4rev1"
	if !s.tlsActive {
		caps += " STARTTLS LOGINDISABLED"
	} else {
		caps += " AUTH=LOGIN"
	}
	s.untagged(caps)
	s.tagged(tag, "OK CAPABILITY completed")
}

func (s *session) cmdStartTLS(tag string) {
	if s.tlsActive {
		s.tagged(tag, "BAD TLS already active")
		return
	}
	s.tagged(tag, "OK begin TLS negotiation now")
	if err := s.upgradeTLS(s.server.tlsConf); err != nil {
		return // connection is likely unusable now; caller's read loop will error out and close
	}
	s.tlsActive = true
}

func (s *session) cmdLogin(tag string, args []string) {
	if !s.tlsActive {
		s.tagged(tag, "NO LOGIN over plaintext refused — use STARTTLS or connect on the implicit-TLS port")
		return
	}

	// Checked before attempting any credential verification — same
	// rationale as smtp.session.handleAuth's authLimiter check.
	ip := connHost(s.conn.RemoteAddr())
	if !s.server.authLimiter.Allow(ip) {
		s.tagged(tag, "NO too many authentication attempts, try again later")
		return
	}

	if len(args) < 2 {
		s.tagged(tag, "BAD LOGIN requires username and password")
		return
	}
	username, password := args[0], args[1]

	if !s.authenticateUser(username, password) {
		s.tagged(tag, "NO LOGIN failed")
		return
	}
	s.tagged(tag, "OK LOGIN completed")
}

func (s *session) cmdSelectExamine(tag string, args []string, readWrite bool) {
	if !s.requireAuthenticated(tag) {
		return
	}
	if len(args) < 1 {
		s.tagged(tag, "BAD SELECT/EXAMINE requires a mailbox name")
		return
	}
	mailbox := args[0]

	entries, err := s.server.database.ListMailboxEntries(s.user.ID, mailbox)
	if err != nil {
		s.tagged(tag, "NO SELECT failed: "+err.Error())
		return
	}

	s.mailbox = mailbox
	s.entries = entries
	s.readOnly = !readWrite
	s.state = stateSelected

	unseen := 0
	nextUID := 1
	for i, e := range entries {
		if !strings.Contains(e.Flags, "\\Seen") && unseen == 0 {
			s.untagged(fmt.Sprintf("OK [UNSEEN %d] first unseen", i+1))
			unseen = i + 1
		}
		if e.UID >= nextUID {
			nextUID = e.UID + 1
		}
	}

	s.untagged(fmt.Sprintf("%d EXISTS", len(entries)))
	s.untagged("0 RECENT")
	s.untagged("FLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)")
	s.untagged("OK [PERMANENTFLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)] Limited")
	s.untagged("OK [UIDVALIDITY 1] UIDs valid")
	s.untagged(fmt.Sprintf("OK [UIDNEXT %d] Predicted next UID", nextUID))

	if readWrite {
		s.tagged(tag, "OK [READ-WRITE] SELECT completed")
	} else {
		s.tagged(tag, "OK [READ-ONLY] EXAMINE completed")
	}
}

func (s *session) cmdList(tag string, args []string) {
	if !s.requireAuthenticated(tag) {
		return
	}
	// args: reference-name mailbox-pattern — we ignore hierarchy and just
	// list every mailbox the user has, since GoMail's folder model is flat
	// (no nested folders yet). A "%"/"*" wildcard pattern matches everything
	// in this simplified model.
	names, err := s.server.database.ListMailboxNames(s.user.ID)
	if err != nil {
		s.tagged(tag, "NO LIST failed: "+err.Error())
		return
	}
	for _, name := range names {
		s.untagged(fmt.Sprintf(`LIST () "/" %s`, quoteIfNeeded(name)))
	}
	s.tagged(tag, "OK LIST completed")
}

func (s *session) cmdClose(tag string) {
	if !s.requireSelected(tag) {
		return
	}
	s.expungeDeleted()
	s.mailbox = ""
	s.entries = nil
	s.state = stateAuthenticated
	s.tagged(tag, "OK CLOSE completed")
}

func (s *session) cmdExpunge(tag string) {
	if !s.requireSelected(tag) {
		return
	}
	if s.readOnly {
		s.tagged(tag, "NO mailbox is read-only")
		return
	}
	removed := s.expungeDeleted()
	s.tagged(tag, fmt.Sprintf("OK EXPUNGE completed (%d removed)", removed))
}

// expungeDeleted removes every \Deleted-flagged message from storage and the
// index, sends the required untagged "N EXPUNGE" responses (in descending
// sequence order, per RFC 3501 §6.4.3 — removing from the end first keeps
// earlier sequence numbers stable for any remaining EXPUNGE responses in the
// same batch), and refreshes the in-memory snapshot.
func (s *session) expungeDeleted() int {
	var kept []db.MailboxEntry
	var removedSeqs []int

	for i, e := range s.entries {
		if strings.Contains(e.Flags, "\\Deleted") {
			removedSeqs = append(removedSeqs, i+1)
			s.server.database.DeleteMailboxEntry(e.ID)
			// Best-effort file removal — the DB row is the source of truth for
			// "does this message exist"; a leftover encrypted file with no
			// index row is inert.
		} else {
			kept = append(kept, e)
		}
	}

	for i := len(removedSeqs) - 1; i >= 0; i-- {
		s.untagged(fmt.Sprintf("%d EXPUNGE", removedSeqs[i]))
	}

	s.entries = kept
	return len(removedSeqs)
}

func (s *session) cmdUID(tag string, args []string) {
	if len(args) < 1 {
		s.tagged(tag, "BAD UID requires a subcommand")
		return
	}
	sub := strings.ToUpper(args[0])
	rest := args[1:]

	switch sub {
	case "FETCH":
		s.cmdFetch(tag, rest, true)
	case "STORE":
		s.cmdStore(tag, rest, true)
	case "SEARCH":
		s.cmdSearch(tag, rest, true)
	default:
		s.tagged(tag, "BAD UID subcommand not recognized")
	}
}

// ── FETCH ─────────────────────────────────────────────────────────────────────

func (s *session) cmdFetch(tag string, args []string, byUID bool) {
	if !s.requireSelected(tag) {
		return
	}
	if len(args) < 2 {
		s.tagged(tag, "BAD FETCH requires a sequence-set and item list")
		return
	}

	targets := s.resolveSequenceSet(args[0], byUID)
	items := expandFetchItems(args[1])

	for _, idx := range targets {
		entry := s.entries[idx]
		s.sendFetchResponse(idx+1, entry, items)
	}
	s.tagged(tag, "OK FETCH completed")
}

func expandFetchItems(token string) []string {
	var items []string
	if isList(token) {
		items = splitList(token)
	} else {
		items = []string{token}
	}
	var expanded []string
	for _, item := range items {
		switch strings.ToUpper(item) {
		case "FAST":
			expanded = append(expanded, "FLAGS", "INTERNALDATE", "RFC822.SIZE")
		case "ALL":
			expanded = append(expanded, "FLAGS", "INTERNALDATE", "RFC822.SIZE")
		case "FULL":
			expanded = append(expanded, "FLAGS", "INTERNALDATE", "RFC822.SIZE", "BODY[]")
		default:
			expanded = append(expanded, item)
		}
	}
	return expanded
}

func (s *session) sendFetchResponse(seq int, entry db.MailboxEntry, items []string) {
	var parts []string
	markSeen := false

	for _, item := range items {
		upper := strings.ToUpper(item)
		switch {
		case upper == "FLAGS":
			parts = append(parts, "FLAGS ("+flagsToIMAP(entry.Flags)+")")
		case upper == "UID":
			parts = append(parts, fmt.Sprintf("UID %d", entry.UID))
		case upper == "RFC822.SIZE":
			parts = append(parts, fmt.Sprintf("RFC822.SIZE %d", entry.SizeBytes))
		case upper == "INTERNALDATE":
			parts = append(parts, fmt.Sprintf(`INTERNALDATE "%s"`, entry.InternalDate.Format("02-Jan-2006 15:04:05 -0700")))
		case upper == "BODY[]" || upper == "RFC822":
			raw, err := s.server.store.Read(entry.EMLPath)
			if err == nil {
				parts = append(parts, fmt.Sprintf("BODY[] {%d}\r\n%s", len(raw), raw))
				markSeen = true
			}
		case upper == "BODY.PEEK[]":
			raw, err := s.server.store.Read(entry.EMLPath)
			if err == nil {
				parts = append(parts, fmt.Sprintf("BODY[] {%d}\r\n%s", len(raw), raw))
			}
		case upper == "BODY[HEADER]" || upper == "RFC822.HEADER" || upper == "BODY.PEEK[HEADER]":
			raw, err := s.server.store.Read(entry.EMLPath)
			if err == nil {
				headers := extractHeaders(raw)
				parts = append(parts, fmt.Sprintf("BODY[HEADER] {%d}\r\n%s", len(headers), headers))
				if upper == "RFC822.HEADER" {
					markSeen = true
				}
			}
		}
	}

	if markSeen && !strings.Contains(entry.Flags, "\\Seen") {
		newFlags := addFlag(entry.Flags, "\\Seen")
		s.server.database.UpdateMailboxFlags(entry.ID, newFlags)
		for i := range s.entries {
			if s.entries[i].ID == entry.ID {
				s.entries[i].Flags = newFlags
			}
		}
	}

	s.untagged(fmt.Sprintf("%d FETCH (%s)", seq, strings.Join(parts, " ")))
}

func extractHeaders(raw []byte) []byte {
	sep := []byte("\r\n\r\n")
	if idx := indexOf(raw, sep); idx >= 0 {
		return raw[:idx+2]
	}
	return raw
}

func indexOf(haystack, needle []byte) int {
	for i := 0; i+len(needle) <= len(haystack); i++ {
		match := true
		for j := range needle {
			if haystack[i+j] != needle[j] {
				match = false
				break
			}
		}
		if match {
			return i
		}
	}
	return -1
}

// ── STORE ─────────────────────────────────────────────────────────────────────

func (s *session) cmdStore(tag string, args []string, byUID bool) {
	if !s.requireSelected(tag) {
		return
	}
	if s.readOnly {
		s.tagged(tag, "NO mailbox is read-only")
		return
	}
	if len(args) < 3 {
		s.tagged(tag, "BAD STORE requires sequence-set, item, and flag list")
		return
	}

	targets := s.resolveSequenceSet(args[0], byUID)
	action := strings.ToUpper(args[1])
	newFlags := splitList(args[2])
	if len(newFlags) == 0 {
		newFlags = args[2:]
	}

	silent := strings.Contains(action, ".SILENT")

	for _, idx := range targets {
		entry := &s.entries[idx]
		switch {
		case strings.HasPrefix(action, "+FLAGS"):
			for _, f := range newFlags {
				entry.Flags = addFlag(entry.Flags, f)
			}
		case strings.HasPrefix(action, "-FLAGS"):
			for _, f := range newFlags {
				entry.Flags = removeFlag(entry.Flags, f)
			}
		case strings.HasPrefix(action, "FLAGS"):
			entry.Flags = strings.Join(newFlags, " ")
		default:
			continue
		}
		s.server.database.UpdateMailboxFlags(entry.ID, entry.Flags)

		if !silent {
			s.untagged(fmt.Sprintf("%d FETCH (FLAGS (%s))", idx+1, flagsToIMAP(entry.Flags)))
		}
	}

	s.tagged(tag, "OK STORE completed")
}

func addFlag(flags, flag string) string {
	if strings.Contains(flags, flag) {
		return flags
	}
	if flags == "" {
		return flag
	}
	return flags + " " + flag
}

func removeFlag(flags, flag string) string {
	parts := strings.Fields(flags)
	var out []string
	for _, p := range parts {
		if p != flag {
			out = append(out, p)
		}
	}
	return strings.Join(out, " ")
}

func flagsToIMAP(flags string) string {
	return flags // stored representation already matches IMAP flag syntax
}

// ── SEARCH ────────────────────────────────────────────────────────────────────

func (s *session) cmdSearch(tag string, args []string, byUID bool) {
	if !s.requireSelected(tag) {
		return
	}
	if len(args) == 0 {
		s.tagged(tag, "BAD SEARCH requires criteria")
		return
	}

	var matches []int
	for i, entry := range s.entries {
		if matchesSearch(entry, args) {
			if byUID {
				matches = append(matches, entry.UID)
			} else {
				matches = append(matches, i+1)
			}
		}
	}

	strs := make([]string, len(matches))
	for i, m := range matches {
		strs[i] = strconv.Itoa(m)
	}
	s.untagged("SEARCH " + strings.Join(strs, " "))
	s.tagged(tag, "OK SEARCH completed")
}

// matchesSearch supports a pragmatic subset: ALL, UNSEEN, SEEN, ANSWERED,
// DELETED, FLAGGED, plus one-shot FROM/SUBJECT substring matching (checked
// against the flags string / a lightweight header scan). Full IMAP SEARCH
// grammar (nested boolean groups, date ranges, OR) is deferred.
func matchesSearch(entry db.MailboxEntry, criteria []string) bool {
	for i := 0; i < len(criteria); i++ {
		switch strings.ToUpper(criteria[i]) {
		case "ALL":
			continue
		case "UNSEEN":
			if strings.Contains(entry.Flags, "\\Seen") {
				return false
			}
		case "SEEN":
			if !strings.Contains(entry.Flags, "\\Seen") {
				return false
			}
		case "ANSWERED":
			if !strings.Contains(entry.Flags, "\\Answered") {
				return false
			}
		case "DELETED":
			if !strings.Contains(entry.Flags, "\\Deleted") {
				return false
			}
		case "FLAGGED":
			if !strings.Contains(entry.Flags, "\\Flagged") {
				return false
			}
		}
	}
	return true
}

// ── Sequence set resolution ────────────────────────────────────────────────────

// resolveSequenceSet parses "1", "1:3", "1,3,5", "1:*" (sequence numbers) or
// the equivalent for UIDs when byUID is true, and returns 0-based indexes
// into s.entries.
func (s *session) resolveSequenceSet(spec string, byUID bool) []int {
	var result []int
	seen := map[int]bool{}

	for _, part := range strings.Split(spec, ",") {
		var lo, hi int
		if strings.Contains(part, ":") {
			bounds := strings.SplitN(part, ":", 2)
			lo = parseSeqNum(bounds[0], byUID, s.entries)
			hi = parseSeqNum(bounds[1], byUID, s.entries)
			if lo > hi {
				lo, hi = hi, lo
			}
		} else {
			lo = parseSeqNum(part, byUID, s.entries)
			hi = lo
		}

		for i, e := range s.entries {
			var val int
			if byUID {
				val = e.UID
			} else {
				val = i + 1
			}
			if val >= lo && val <= hi && !seen[i] {
				seen[i] = true
				result = append(result, i)
			}
		}
	}
	return result
}

func parseSeqNum(s string, byUID bool, entries []db.MailboxEntry) int {
	if s == "*" {
		if len(entries) == 0 {
			return 0
		}
		if byUID {
			return entries[len(entries)-1].UID
		}
		return len(entries)
	}
	n, err := strconv.Atoi(s)
	if err != nil {
		return 0
	}
	return n
}

func quoteIfNeeded(name string) string {
	if strings.ContainsAny(name, " \t()\"") {
		return `"` + strings.ReplaceAll(name, `"`, `\"`) + `"`
	}
	return name
}