mailgoserver/internal/webui/dkim.go

package webui

import (
	"encoding/json"
	"net/http"
	"regexp"
	"strings"
	"time"

	"mailgoserver/internal/db"
	"mailgoserver/internal/dkim"
)

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(v)
}

func isAjax(r *http.Request) bool {
	return r.Header.Get("X-Requested-With") == "XMLHttpRequest" || strings.Contains(r.Header.Get("Content-Type"), "application/json")
}

// dkimList mirrors dkim.py's dkim_list().
func (a *App) dkimList(w http.ResponseWriter, r *http.Request) {
	active, err := a.DB.ListActiveDKIMKeysWithDomain()
	if err != nil {
		setFlash(w, "error", "Error loading DKIM keys")
	}
	inactive, _ := a.DB.ListInactiveDKIMKeysWithDomain()
	publicIP := getPublicIP(a.Cfg)
	scope := scopeFromContext(r)

	var dkimData []M
	for _, k := range active {
		if !scope.Allowed(k.DomainID) {
			continue
		}
		rec, _ := a.DKIM.GetDKIMPublicKeyRecord(k.DomainName)
		if rec == nil {
			rec = &dkim.DNSRecord{}
		}
		spfCheck := checkDNSRecord(k.DomainName)
		existingSPF := ""
		for _, txt := range spfCheck.Records {
			if strings.Contains(txt, "v=spf1") {
				existingSPF = txt
				break
			}
		}
		dkimData = append(dkimData, M{
			"domain":          M{"id": k.DomainID, "domain_name": k.DomainName},
			"dkim_key":        k.DKIMKey,
			"dns_record":      M{"name": rec.Name, "value": rec.Value},
			"existing_spf":    existingSPF,
			"recommended_spf": generateSPFRecord(publicIP, existingSPF),
			"public_ip":       publicIP,
		})
	}
	var oldData []M
	for _, k := range inactive {
		if !scope.Allowed(k.DomainID) {
			continue
		}
		status := "Disabled"
		if k.ReplacedAt != nil {
			status = "Replaced"
		}
		oldData = append(oldData, M{
			"domain":      M{"id": k.DomainID, "domain_name": k.DomainName},
			"dkim_key":    k.DKIMKey,
			"status_text": status,
		})
	}

	a.render(w, r, "dkim.html", M{"active": "dkim", "dkim_data": dkimData, "old_dkim_data": oldData})
}

// createDKIM mirrors dkim.py's create_dkim().
func (a *App) createDKIM(w http.ResponseWriter, r *http.Request) {
	domain := r.FormValue("domain")
	selector := r.FormValue("selector")
	if domain == "" {
		writeJSON(w, http.StatusBadRequest, M{"success": false, "message": "Domain is required"})
		return
	}
	dom, err := a.DB.GetDomainByNameExact(domain)
	if err != nil || dom == nil {
		writeJSON(w, http.StatusNotFound, M{"success": false, "message": "Domain not found"})
		return
	}
	if !scopeFromContext(r).Allowed(dom.ID) {
		writeJSON(w, http.StatusNotFound, M{"success": false, "message": "Domain not found"})
		return
	}
	if err := a.DB.DeactivateActiveDKIMKeysForDomain(dom.ID, time.Now()); err != nil {
		writeJSON(w, http.StatusInternalServerError, M{"success": false, "message": "Failed to create DKIM key"})
		return
	}
	ok, err := a.DKIM.GenerateDKIMKeypair(domain, selector, true)
	if err != nil || !ok {
		writeJSON(w, http.StatusInternalServerError, M{"success": false, "message": "Failed to create DKIM key"})
		return
	}
	writeJSON(w, http.StatusOK, M{"success": true, "message": "DKIM key created successfully"})
}

// regenerateDKIM mirrors dkim.py's regenerate_dkim() — path id is a DOMAIN id.
func (a *App) regenerateDKIM(w http.ResponseWriter, r *http.Request) {
	domainID := pathID(r)
	dom, err := a.DB.GetDomainByID(domainID)
	if err != nil || dom == nil {
		http.NotFound(w, r)
		return
	}
	if !requireDomainAccess(w, r, dom.ID) {
		return
	}
	current, _ := a.DB.GetActiveDKIMKeyByDomainID(domainID)
	selector := ""
	if current != nil {
		selector = current.Selector
	}
	if err := a.DB.DeactivateActiveDKIMKeysForDomain(domainID, time.Now()); err != nil {
		a.dkimActionFailed(w, r, "Failed to regenerate DKIM key")
		return
	}
	ok, err := a.DKIM.GenerateDKIMKeypair(dom.DomainName, selector, true)
	if err != nil || !ok {
		a.dkimActionFailed(w, r, "Failed to regenerate DKIM key")
		return
	}

	if isAjax(r) {
		publicIP := getPublicIP(a.Cfg)
		rec, _ := a.DKIM.GetDKIMPublicKeyRecord(dom.DomainName)
		spfCheck := checkDNSRecord(dom.DomainName)
		existingSPF := ""
		for _, txt := range spfCheck.Records {
			if strings.Contains(txt, "v=spf1") {
				existingSPF = txt
				break
			}
		}
		newKey, _ := a.DB.GetActiveDKIMKeyByDomainID(domainID)
		writeJSON(w, http.StatusOK, M{
			"success": true, "message": "DKIM key regenerated successfully",
			"new_key": newKey, "dns_record": rec, "existing_spf": existingSPF,
			"recommended_spf": generateSPFRecord(publicIP, existingSPF),
			"public_ip":       publicIP, "domain": dom.DomainName,
		})
		return
	}
	setFlash(w, "success", "DKIM key regenerated successfully")
	http.Redirect(w, r, Prefix+"/dkim", http.StatusFound)
}

func (a *App) dkimActionFailed(w http.ResponseWriter, r *http.Request, message string) {
	if isAjax(r) {
		writeJSON(w, http.StatusInternalServerError, M{"success": false, "message": message})
		return
	}
	setFlash(w, "error", message)
	http.Redirect(w, r, Prefix+"/dkim", http.StatusFound)
}

func (a *App) editDKIMForm(w http.ResponseWriter, r *http.Request) {
	key, ok := a.dkimKeyWithAccess(w, r)
	if !ok {
		return
	}
	dom, _ := a.DB.GetDomainByID(key.DomainID)
	a.render(w, r, "edit_dkim.html", M{"active": "dkim", "dkim_key": key, "domain": dom})
}

var selectorPattern = regexp.MustCompile(`^[a-zA-Z0-9_-]+$`)

// dkimKeyWithAccess fetches a DKIM key by path ID and confirms it belongs to a domain
// the current admin can manage.
func (a *App) dkimKeyWithAccess(w http.ResponseWriter, r *http.Request) (key *db.DKIMKey, ok bool) {
	key, err := a.DB.GetDKIMKeyByID(pathID(r))
	if err != nil || key == nil {
		http.NotFound(w, r)
		return nil, false
	}
	if !requireDomainAccess(w, r, key.DomainID) {
		return nil, false
	}
	return key, true
}

// editDKIM mirrors dkim.py's edit_dkim() POST branch.
func (a *App) editDKIM(w http.ResponseWriter, r *http.Request) {
	key, ok := a.dkimKeyWithAccess(w, r)
	if !ok {
		return
	}
	id := key.ID
	dom, _ := a.DB.GetDomainByID(key.DomainID)
	selector := strings.TrimSpace(r.FormValue("selector"))

	if selector == "" || !selectorPattern.MatchString(selector) {
		setFlash(w, "error", "Selector must contain only letters, numbers, hyphens, and underscores")
		a.render(w, r, "edit_dkim.html", M{"active": "dkim", "dkim_key": key, "domain": dom})
		return
	}
	if exists, _ := a.DB.SelectorExistsForDomain(key.DomainID, selector, id); exists {
		setFlash(w, "error", "This selector is already in use for this domain")
		a.render(w, r, "edit_dkim.html", M{"active": "dkim", "dkim_key": key, "domain": dom})
		return
	}
	if err := a.DB.UpdateDKIMKeySelector(id, selector); err != nil {
		setFlash(w, "error", "Error updating selector")
		http.Redirect(w, r, Prefix+"/dkim", http.StatusFound)
		return
	}
	setFlash(w, "success", "DKIM selector updated successfully")
	http.Redirect(w, r, Prefix+"/dkim", http.StatusFound)
}

// toggleDKIM mirrors dkim.py's toggle_dkim().
func (a *App) toggleDKIM(w http.ResponseWriter, r *http.Request) {
	key, ok := a.dkimKeyWithAccess(w, r)
	if !ok {
		return
	}
	id := key.ID
	newActive := !key.IsActive
	if newActive {
		if err := a.DB.DeactivateActiveDKIMKeysForDomain(key.DomainID, time.Now()); err != nil {
			a.dkimActionFailed(w, r, "Error toggling DKIM status")
			return
		}
	}
	if err := a.DB.SetDKIMKeyActive(id, newActive, time.Now()); err != nil {
		a.dkimActionFailed(w, r, "Error toggling DKIM status")
		return
	}
	message := "DKIM key disabled"
	if newActive {
		message = "DKIM key enabled"
	}
	if isAjax(r) {
		writeJSON(w, http.StatusOK, M{"success": true, "message": message, "is_active": newActive})
		return
	}
	setFlash(w, "success", message)
	http.Redirect(w, r, Prefix+"/dkim", http.StatusFound)
}

func (a *App) removeDKIM(w http.ResponseWriter, r *http.Request) {
	key, ok := a.dkimKeyWithAccess(w, r)
	if !ok {
		return
	}
	if err := a.DB.RemoveDKIMKey(key.ID); err != nil {
		setFlash(w, "error", "Error removing DKIM key")
	} else {
		setFlash(w, "success", "DKIM key permanently removed")
	}
	http.Redirect(w, r, Prefix+"/dkim", http.StatusFound)
}

// checkDKIMDNS mirrors dkim.py's check_dkim_dns().
func (a *App) checkDKIMDNS(w http.ResponseWriter, r *http.Request) {
	domain := r.FormValue("domain")
	selector := r.FormValue("selector")
	rec, err := a.DKIM.GetDKIMPublicKeyRecord(domain)
	if err != nil || rec == nil {
		writeJSON(w, http.StatusOK, M{"success": false, "message": "No active DKIM key for domain"})
		return
	}
	dnsName := selector + "._domainkey." + domain
	result := checkDNSRecord(dnsName)
	found := false
	expected := strings.Trim(rec.Value, `"`)
	for _, txt := range result.Records {
		if strings.Contains(txt, expected) || strings.Contains(expected, txt) {
			found = true
			break
		}
	}
	message := "DKIM record not found or does not match"
	if found {
		message = "DKIM record found and matches"
	} else if !result.Success {
		message = result.Message
	}
	writeJSON(w, http.StatusOK, M{"success": found, "message": message, "records": result.Records})
}

// checkSPFDNS mirrors dkim.py's check_spf_dns().
func (a *App) checkSPFDNS(w http.ResponseWriter, r *http.Request) {
	domain := r.FormValue("domain")
	result := checkDNSRecord(domain)
	var spfRecord string
	for _, txt := range result.Records {
		if strings.Contains(txt, "v=spf1") {
			spfRecord = txt
			break
		}
	}
	publicIP := getPublicIP(a.Cfg)
	validForServer := spfRecord != "" && strings.Contains(spfRecord, "ip4:"+publicIP)
	message := "SPF record not found"
	if spfRecord != "" {
		message = "SPF record found"
	}
	writeJSON(w, http.StatusOK, M{
		"success": spfRecord != "", "message": message, "records": result.Records,
		"spf_record": spfRecord, "spf_valid_for_server": validForServer,
		"spf_check_message": message, "public_ip": publicIP,
	})
}